@vellumai/assistant 0.8.4 → 0.8.5

package/src/memory/v3/shadow-middleware.ts

@@ -0,0 +1,305 @@
+/**
+ * Memory v3 — live-shadow `memoryRetrieval` middleware.
+ *
+ * Registered unconditionally into the `memoryRetrieval` pipeline, but inert
+ * unless BOTH `config.memory.v3.enabled` and `config.memory.v3.shadow` are on.
+ * When inert it is a byte-for-byte pass-through: it returns `next(args)`
+ * verbatim and performs zero extra work (no v3 call, no DB read, no log write).
+ *
+ * When active, it:
+ *   1. Returns the real (v2/default) `MemoryResult` from `next(args)` promptly —
+ *      the injected context is ALWAYS the v2 result, never v3.
+ *   2. Kicks off the v3 retrieval loop DETACHED (not awaited on the path that
+ *      returns the result), so the shadow run can never block or slow the turn.
+ *   3. Logs v3's selection set to `memory_v2_activation_logs` with
+ *      `mode = "v3_shadow"`. The harness oracle filters `mode='router'`, so
+ *      shadow rows never pollute it; the inspector can still surface them.
+ *
+ * The shadow build mirrors the inputs the v2 router receives (recent turn
+ * pairs, NOW context, prior-ever-injected slugs, config) so its recall is
+ * measured against the same situational context the live path saw. Failures
+ * are swallowed with a warn — the shadow is observational only and must never
+ * affect the live turn.
+ */
+
+import { desc, eq } from "drizzle-orm";
+
+import { getConfig } from "../../config/loader.js";
+import { registerPlugin } from "../../plugins/registry.js";
+import {
+  type MemoryArgs,
+  type MemoryResult,
+  type Middleware,
+  type Plugin,
+  PluginExecutionError,
+} from "../../plugins/types.js";
+import type { ContentBlock } from "../../providers/types.js";
+import { getLogger } from "../../util/logger.js";
+import { getWorkspaceDir } from "../../util/platform.js";
+import type { DrizzleDb } from "../db-connection.js";
+import { getDb } from "../db-connection.js";
+import {
+  type MemoryV2ConceptRowRecord,
+  type MemoryV2ConfigSnapshot,
+  recordMemoryV2ActivationLog,
+} from "../memory-v2-activation-log-store.js";
+import { messages } from "../schema.js";
+import { hydrate } from "../v2/activation-store.js";
+import type { RetrievalInput } from "../v2/harness/retriever.js";
+import { loadNowText } from "../v2/now-text.js";
+import type { RouterTurnPair } from "../v2/router.js";
+import type { EverInjectedEntry } from "../v2/types.js";
+import { runRetrievalLoop } from "./loop.js";
+
+const log = getLogger("memory-v3-shadow");
+
+/**
+ * Extract the recent (assistant, user) turn pairs from a conversation's
+ * message list, newest-pair-last, capped at `k`. Mirrors production
+ * `extractRecentTurnPairs` in `conversation-graph-memory.ts` (and its harness
+ * twin in `replay-input.ts`) so the shadow's `recentTurnPairs` matches what the
+ * live router was fed.
+ */
+function extractRecentTurnPairs(
+  msgs: ReadonlyArray<{ role: string; content: ContentBlock[] }>,
+  k: number,
+): RouterTurnPair[] {
+  const messageText = (content: ContentBlock[]): string =>
+    content
+      .filter(
+        (b): b is Extract<ContentBlock, { type: "text" }> => b.type === "text",
+      )
+      .map((b) => b.text)
+      .join(" ");
+
+  const pairs: RouterTurnPair[] = [];
+  let pendingUser: string | null = null;
+  for (let i = msgs.length - 1; i >= 0 && pairs.length < k; i--) {
+    const msg = msgs[i]!;
+    if (msg.role === "user" && pendingUser === null) {
+      pendingUser = messageText(msg.content);
+    } else if (msg.role === "assistant" && pendingUser !== null) {
+      pairs.unshift({
+        assistantMessage: messageText(msg.content),
+        userMessage: pendingUser,
+      });
+      pendingUser = null;
+    }
+  }
+  if (pendingUser !== null && pairs.length < k) {
+    pairs.unshift({ assistantMessage: "", userMessage: pendingUser });
+  }
+  if (pairs.length === 0) {
+    pairs.push({ assistantMessage: "", userMessage: "" });
+  }
+  return pairs;
+}
+
+/** Parse a persisted JSON content-block string; tolerate malformed rows. */
+function parseContent(raw: string): ContentBlock[] {
+  try {
+    const parsed = JSON.parse(raw);
+    return Array.isArray(parsed) ? (parsed as ContentBlock[]) : [];
+  } catch {
+    return [];
+  }
+}
+
+/**
+ * Load the most recent messages for a conversation, oldest-first, bounded to a
+ * small generous multiple of `historicalPairs`. Pair extraction only needs the
+ * tail, so a bounded `LIMIT` query avoids loading an entire (potentially
+ * multi-GB) conversation on every shadow turn — mirrors the harness's bounded
+ * fetch in `replay-input.ts`.
+ */
+function loadRecentMessages(
+  db: DrizzleDb,
+  conversationId: string,
+  historicalPairs: number,
+): Array<{ role: string; content: ContentBlock[] }> {
+  const fetchWindow = Math.max(20, historicalPairs * 12);
+  const rows = db
+    .select({ role: messages.role, content: messages.content })
+    .from(messages)
+    .where(eq(messages.conversationId, conversationId))
+    .orderBy(desc(messages.createdAt), desc(messages.id))
+    .limit(fetchWindow)
+    .all();
+  return rows
+    .reverse()
+    .map((r) => ({ role: r.role, content: parseContent(r.content) }));
+}
+
+/**
+ * Empty config snapshot for shadow log rows. The activation-state values are
+ * meaningless for a v3 selection (it computes no spreading-activation scores),
+ * so they are zeroed — exactly as the v2 router-mode rows do.
+ */
+const SHADOW_CONFIG_SNAPSHOT: MemoryV2ConfigSnapshot = {
+  d: 0,
+  c_user: 0,
+  c_assistant: 0,
+  c_now: 0,
+  k: 0,
+  hops: 0,
+  top_k: 0,
+  epsilon: 0,
+};
+
+/**
+ * Build the concept rows logged for a v3 shadow selection. Each selected slug
+ * becomes a zeroed concept row tagged `source: "router"` and
+ * `status: "injected"` — the shadow has no activation scores to record, and the
+ * `mode='v3_shadow'` row tag (not the concept source) is what distinguishes
+ * shadow telemetry from live router selections.
+ */
+function buildShadowConceptRows(
+  selectedSlugs: readonly string[],
+): MemoryV2ConceptRowRecord[] {
+  return selectedSlugs.map((slug) => ({
+    slug,
+    finalActivation: 0,
+    ownActivation: 0,
+    priorActivation: 0,
+    simUser: 0,
+    simAssistant: 0,
+    simNow: 0,
+    simUserRerankBoost: 0,
+    simAssistantRerankBoost: 0,
+    inRerankPool: false,
+    spreadContribution: 0,
+    source: "router",
+    status: "injected",
+  }));
+}
+
+/**
+ * Run the v3 retrieval loop for the shadow and log its selection. Best-effort:
+ * any failure is logged and swallowed. Honors `signal` so a cancelled turn
+ * stops the shadow's lane work.
+ */
+async function runShadowAndLog(
+  args: MemoryArgs,
+  signal: AbortSignal,
+): Promise<void> {
+  try {
+    if (signal.aborted) return;
+
+    const config = getConfig();
+    const workspaceDir = getWorkspaceDir();
+    const db = getDb();
+
+    const historicalPairs = config.memory.v2.router.historical_pairs;
+    const recentMessages = loadRecentMessages(
+      db,
+      args.conversationId,
+      historicalPairs,
+    );
+    const recentTurnPairs = extractRecentTurnPairs(
+      recentMessages,
+      historicalPairs,
+    );
+
+    const nowText = await loadNowText(workspaceDir);
+
+    let priorEverInjected: readonly EverInjectedEntry[] = [];
+    try {
+      const state = await hydrate(db, args.conversationId);
+      priorEverInjected = state?.everInjected ?? [];
+    } catch (err) {
+      log.warn(
+        { err, conversationId: args.conversationId },
+        "v3 shadow: failed to hydrate prior-ever-injected; continuing with empty set",
+      );
+    }
+
+    if (signal.aborted) return;
+
+    const input: RetrievalInput = {
+      workspaceDir,
+      recentTurnPairs,
+      nowText,
+      priorEverInjected,
+      config,
+      signal,
+    };
+
+    const output = await runRetrievalLoop(input, {
+      db,
+      conversationId: args.conversationId,
+      turn: args.turnIndex,
+    });
+
+    if (signal.aborted) return;
+
+    recordMemoryV2ActivationLog({
+      conversationId: args.conversationId,
+      turn: args.turnIndex,
+      mode: "v3_shadow",
+      concepts: buildShadowConceptRows(output.selectedSlugs),
+      config: SHADOW_CONFIG_SNAPSHOT,
+    });
+  } catch (err) {
+    log.warn(
+      { err, conversationId: args.conversationId, turn: args.turnIndex },
+      "v3 shadow retrieval failed; live turn unaffected",
+    );
+  }
+}
+
+/**
+ * Live-shadow `memoryRetrieval` middleware.
+ *
+ * Flag-gated INSIDE the middleware (per-turn, live-toggle): when v3 shadow is
+ * off it is a pure pass-through. When on, it fires the v3 loop detached and
+ * returns the unchanged downstream (v2) result immediately.
+ */
+export const memoryV3ShadowMiddleware: Middleware<MemoryArgs, MemoryResult> =
+  async function memoryV3Shadow(args, next) {
+    const v3 = getConfig().memory.v3;
+    if (!v3?.enabled || !v3?.shadow) {
+      // Inert: byte-for-byte pass-through, zero extra work.
+      return next(args);
+    }
+
+    // Detached — never awaited on the path that returns the result, so the
+    // shadow can neither block nor slow the live turn. Errors are swallowed
+    // inside `runShadowAndLog`.
+    void runShadowAndLog(args, args.signal);
+
+    return next(args);
+  };
+
+/**
+ * First-party plugin contributing the live-shadow `memoryRetrieval`
+ * middleware. Registered unconditionally by the plugin bootstrap (it is inert
+ * unless both v3 flags are on), so the registration is always present but does
+ * zero work in the default (flags-off) configuration.
+ */
+export const memoryV3ShadowPlugin: Plugin = {
+  manifest: {
+    name: "memory-v3-shadow",
+    version: "0.0.1",
+  },
+  middleware: {
+    memoryRetrieval: memoryV3ShadowMiddleware,
+  },
+};
+
+// Module-load side effect: register the shadow plugin at import time so the
+// registry is populated even in tests that skip `bootstrapPlugins()`, matching
+// the first-party `default-*` plugins. Idempotent via the swallowed
+// duplicate-name check (the defaults aggregator also lists this plugin).
+try {
+  registerPlugin(memoryV3ShadowPlugin);
+} catch (err) {
+  if (
+    err instanceof PluginExecutionError &&
+    err.message.includes("already registered")
+  ) {
+    // already registered — expected when both the defaults aggregator and the
+    // direct module import run in the same process.
+  } else {
+    throw err;
+  }
+}

package/src/memory/v3/traversal.ts

@@ -0,0 +1,206 @@
+/**
+ * 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`).
+ *   - The `page:` children the `descend` decision chooses to *keep* are collected
+ *     into the returned `pages` set — pages are leaves, never recursed into. A
+ *     page the decision does not keep is dropped, so the walk emits a curated
+ *     selection rather than every page it passes.
+ *   - A `visited` set keyed by canonical id (`node:<id>`) 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 decision injected into {@link walkTree}. Given a node id and its ordered
+ * child refs, return the *node* children to recurse into and the *page* children
+ * to keep for the answer. The driver 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 `descend` refs that are not `node:` children of `nodeId`, and
+ * `keep` refs that are not `page:` children of `nodeId`, are ignored by the walk
+ * (it only acts on the distinct children it actually offered).
+ */
+export type DescendDecision = (
+  nodeId: string,
+  children: ReadonlyArray<ChildRef>,
+) => Promise<DescendResult> | DescendResult;
+
+/**
+ * The result of a {@link DescendDecision}. `descend` lists the `node:` children
+ * chosen for recursion; `keep` lists the `page:` children to collect into the
+ * walk's result; `reasoning` is the optional model rationale recorded on the
+ * level.
+ */
+export interface DescendResult {
+  descend: ChildRef[];
+  keep: 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 the descend decision kept across the walk, dedup'd. */
+  pages: Set<string>;
+  /** 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<ChildRef> {
+  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<WalkResult> {
+  const { breadthBudget, maxDepth, descend } = opts;
+  const start = opts.start ?? tree.root;
+
+  const pages = new Set<string>();
+  const levels: TreeLevel[] = [];
+  const visited = new Set<string>();
+
+  // 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 only the page children the decision chose to keep, filtered to
+      // the page children this node actually offered — a decision can't keep a
+      // page the node never presented.
+      const offeredPageRefs = new Set(
+        children.filter((c) => c.kind === "page").map((c) => c.ref),
+      );
+      for (const choice of result.keep) {
+        if (choice.kind === "page" && offeredPageRefs.has(choice.ref)) {
+          pages.add(choice.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<string>();
+      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 };
+}