@vellumai/assistant 0.8.5 → 0.8.7

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

@@ -0,0 +1,379 @@
+/**
+ * Memory v3 — flag-gated shadow/live plugin.
+ *
+ * Registered as an {@link Injector} that runs the v3 orchestrator each turn and
+ * records its selection set to `memory_v3_selections`. Two flags gate its
+ * injection behavior:
+ *
+ *   - `memory-v3-shadow` (and live OFF): observation-only. `produce()` returns
+ *     `null`, so v2 injection is bit-for-bit identical whether the flag is on
+ *     or off — the only difference is the side-effect telemetry write.
+ *   - `memory-v3-live`: live injection. `produce()` additionally renders the
+ *     working-set selection into a `<memory>` block via {@link renderMemoryBlock}
+ *     and returns it as an injection block at v2's dynamic-memory placement
+ *     (`after-memory-prefix`). Selections are still logged.
+ *   - both OFF: `produce()` returns `null` and skips orchestration entirely.
+ *
+ * On each turn (either flag on):
+ *   1. Lazy-init the v3 lanes ONCE across the whole process (leaf tree, core
+ *      set, BM25 needle, carry-forward working set), memoizing the init
+ *      promise so concurrent first turns share a single build.
+ *   2. Build a {@link TurnContext} from the conversation's recent messages.
+ *   3. Run {@link orchestrate} and record its selection set to
+ *      `memory_v3_selections` with a best-effort lane attribution.
+ *
+ * Everything after the flag read is wrapped in try/catch — any failure is
+ * logged and swallowed so it can never affect the live turn. In live mode a
+ * failure returns `null` (no v3 injection); v2 suppression keys off BOTH the
+ * flag AND whether `produce()` actually returned a block, so a v3 failure (or
+ * empty selection) falls back to v2 memory rather than dropping all memory.
+ */
+
+import { isAssistantFeatureFlagEnabled } from "../../config/assistant-feature-flags.js";
+import { getConfig } from "../../config/loader.js";
+import type { AssistantConfig } from "../../config/schema.js";
+import { getMessages } from "../../memory/conversation-crud.js";
+import { getDb, getSqliteFrom } from "../../memory/db-connection.js";
+import { stringifyMessageContent } from "../../memory/message-content.js";
+import {
+  type InjectionBlock,
+  type Injector,
+  type Plugin,
+  type TurnContext as PluginTurnContext,
+} from "../../plugins/types.js";
+import { getLogger } from "../../util/logger.js";
+import { getWorkspaceDir } from "../../util/platform.js";
+import { getPageIndex } from "../v2/page-index.js";
+import { injectCapabilitiesLeaf, isCapabilitySlug } from "./capabilities.js";
+import { loadCore } from "./core.js";
+import type { NeedleIndex } from "./needle.js";
+import { buildNeedleIndex } from "./needle.js";
+import type { OrchestrateResult } from "./orchestrate.js";
+import { orchestrate } from "./orchestrate.js";
+import { renderV3PageContent } from "./page-content.js";
+import { renderMemoryBlock } from "./render-injection.js";
+import { coreSlugs, loadLeafTree, resolveDataDir } from "./tree.js";
+import {
+  type LeafPath,
+  type LeafTree,
+  MEMORY_V3_BLOCK_ID,
+  type SelectionSource,
+  type Slug,
+  type TurnContext,
+} from "./types.js";
+import { WorkingSet } from "./working-set.js";
+
+const MEMORY_V3_SHADOW = "memory-v3-shadow" as const;
+const MEMORY_V3_LIVE = "memory-v3-live" as const;
+
+const log = getLogger("memory-v3-shadow");
+
+/** How many recent messages to fold into the shadow `recentContext` string. */
+const RECENT_CONTEXT_MESSAGES = 6;
+
+/**
+ * The lazily-built, process-lifetime v3 lanes. The working set is stateful
+ * (carry-forward across turns), so it is intentionally shared across every
+ * shadow turn rather than rebuilt per turn — that is the whole point of the
+ * carry-forward lane.
+ */
+export interface ShadowLanes {
+  tree: LeafTree;
+  core: Set<LeafPath>;
+  needle: NeedleIndex;
+  workingSet: WorkingSet;
+}
+
+/**
+ * Memoized init promise. Caching the PROMISE (not the resolved value) means
+ * concurrent first turns all await the same build instead of racing several
+ * `loadLeafTree`/`buildNeedleIndex` passes.
+ */
+let lanesPromise: Promise<ShadowLanes> | null = null;
+
+/**
+ * Drop the memoized lanes so the NEXT `getLanes` rebuilds them from scratch
+ * (fresh topic-tree load + fresh needle index). The rebuild is lazy — this only
+ * clears the cache, so the cost is paid by the next caller, and concurrent
+ * first-callers after the invalidation still share a single build via the
+ * re-memoized promise. Call this whenever the underlying tree or needle sources
+ * change on disk.
+ */
+export function invalidateLanes(): void {
+  lanesPromise = null;
+}
+
+/** Test-only alias for {@link invalidateLanes}. */
+export function resetShadowLanesForTests(): void {
+  invalidateLanes();
+}
+
+/**
+ * Pull a page's summary from the existing v2 page index. Missing summaries
+ * (and any read failure) degrade to "" — the needle and L2 selector treat an
+ * empty summary as "no signal" rather than throwing.
+ */
+async function pageSummary(slug: Slug): Promise<string> {
+  try {
+    const index = await getPageIndex(getWorkspaceDir());
+    return index.bySlug.get(slug)?.summary ?? "";
+  } catch {
+    return "";
+  }
+}
+
+async function initLanes(config: AssistantConfig): Promise<ShadowLanes> {
+  const dataDir = resolveDataDir();
+  const pageIndex = await getPageIndex(getWorkspaceDir());
+  const pageLeaves = new Map<Slug, LeafPath[]>();
+  for (const entry of pageIndex.entries) {
+    pageLeaves.set(entry.slug, entry.leaves);
+  }
+  const tree = await loadLeafTree(dataDir, pageLeaves);
+  const core = await loadCore(dataDir);
+
+  // Always-on synthetic capabilities leaf: skill + CLI-command rows the page
+  // index already carries (with summaries). Injecting them as leaf members puts
+  // them in the needle corpus (`tree.byPage`) and, via `core`, makes L1 always
+  // open the leaf so L2 selects the relevant subset each turn — matching v2's
+  // "always in the pool, router-selected" capability surfacing. Done before
+  // `buildNeedleIndex` so the synthetic members are indexed.
+  const capabilitySlugs = pageIndex.entries
+    .map((entry) => entry.slug)
+    .filter(isCapabilitySlug);
+  injectCapabilitiesLeaf(tree, core, capabilitySlugs);
+
+  const needle = await buildNeedleIndex(tree, pageSummary);
+  const workingSet = new WorkingSet(
+    config.memory.v3.workingSet.maxPages,
+    config.memory.v3.workingSet.evictWindow,
+  );
+
+  return { tree, core, needle, workingSet };
+}
+
+/** Lazy, memoized accessor for the shadow lanes. */
+function getLanes(config: AssistantConfig): Promise<ShadowLanes> {
+  if (!lanesPromise) {
+    lanesPromise = initLanes(config).catch((err) => {
+      // Reset on failure so a transient init error doesn't permanently wedge
+      // the shadow lane — the next turn retries.
+      lanesPromise = null;
+      throw err;
+    });
+  }
+  return lanesPromise;
+}
+
+/**
+ * Build a v3 {@link TurnContext} from the conversation's persisted messages.
+ * `currentMessage` is the latest user message; `recentContext` is the tail of
+ * the recent transcript. Returns `null` when there is no user message to route
+ * on (nothing to shadow this turn).
+ */
+function buildShadowTurn(
+  conversationId: string,
+  turnIndex: number,
+): TurnContext | null {
+  const rows = getMessages(conversationId);
+  if (rows.length === 0) return null;
+
+  let currentMessage = "";
+  for (let i = rows.length - 1; i >= 0; i--) {
+    if (rows[i]!.role === "user") {
+      currentMessage = stringifyMessageContent(rows[i]!.content);
+      if (currentMessage.length > 0) break;
+    }
+  }
+  if (currentMessage.length === 0) return null;
+
+  const recentContext = rows
+    .slice(-RECENT_CONTEXT_MESSAGES)
+    .map((r) => stringifyMessageContent(r.content))
+    .filter((t) => t.length > 0)
+    .join("\n");
+
+  return {
+    conversationId,
+    turnNumber: turnIndex,
+    currentMessage,
+    recentContext,
+  };
+}
+
+interface SelectionRow {
+  slug: Slug;
+  source: SelectionSource;
+  pinned: number;
+}
+
+/**
+ * Map an orchestrate result onto telemetry rows with best-effort lane
+ * attribution.
+ *
+ * - A current selection whose page belongs to a core leaf → `"core+l2"`,
+ *   otherwise `"l1+l2"`.
+ * - A slug in `finalInjection` but NOT re-selected this turn → `"carry-forward"`.
+ *
+ * Precise needle attribution is a documented follow-up: the needle lane only
+ * widens the open set, so a needle-surfaced page that survives L2 selection is
+ * indistinguishable here from an L1-routed one. This coarse mapping is
+ * acceptable for v0 shadow telemetry.
+ */
+function attributeSelections(
+  tree: LeafTree,
+  core: Set<LeafPath>,
+  result: OrchestrateResult,
+): SelectionRow[] {
+  const coreOwnedSlugs = coreSlugs(tree, core);
+
+  const rows: SelectionRow[] = [];
+  const seen = new Set<Slug>();
+  for (const sel of result.currentSelections) {
+    seen.add(sel.slug);
+    rows.push({
+      slug: sel.slug,
+      source: coreOwnedSlugs.has(sel.slug) ? "core+l2" : "l1+l2",
+      pinned: sel.pinned ? 1 : 0,
+    });
+  }
+  for (const slug of result.finalInjection) {
+    if (seen.has(slug)) continue;
+    seen.add(slug);
+    rows.push({ slug, source: "carry-forward", pinned: 0 });
+  }
+  return rows;
+}
+
+/** Write the attributed selection rows to `memory_v3_selections`. */
+function writeSelections(
+  conversationId: string,
+  turn: number,
+  rows: SelectionRow[],
+): void {
+  if (rows.length === 0) return;
+  const raw = getSqliteFrom(getDb());
+  // PK is (conversation_id, turn, slug); OR REPLACE keeps the write
+  // idempotent if the same turn is observed twice (e.g. a retried turn).
+  const stmt = raw.query(/*sql*/ `
+    INSERT OR REPLACE INTO memory_v3_selections (
+      conversation_id, turn, slug, source, pinned, created_at
+    ) VALUES (?, ?, ?, ?, ?, ?)
+  `);
+  const now = Date.now();
+  for (const row of rows) {
+    stmt.run(conversationId, turn, row.slug, row.source, row.pinned, now);
+  }
+}
+
+/**
+ * Run v3 orchestration for one turn and log the selection set, returning the
+ * orchestrate result so a live caller can render it. Never throws — all
+ * failures are logged and swallowed (returning `null`) so the live turn is
+ * unaffected. Returns `null` when there is no user message to route on.
+ */
+async function observeTurn(
+  conversationId: string,
+  turnIndex: number,
+): Promise<OrchestrateResult | null> {
+  try {
+    const turn = buildShadowTurn(conversationId, turnIndex);
+    if (!turn) return null;
+
+    const lanes = await getLanes(getConfig());
+    const result = await orchestrate(turn, {
+      tree: lanes.tree,
+      core: lanes.core,
+      needle: lanes.needle,
+      workingSet: lanes.workingSet,
+      pageSummary,
+    });
+
+    const rows = attributeSelections(lanes.tree, lanes.core, result);
+    writeSelections(conversationId, turnIndex, rows);
+    return result;
+  } catch (err) {
+    log.warn(
+      { err: err instanceof Error ? err.message : String(err), conversationId },
+      "memory-v3 orchestration failed (non-fatal)",
+    );
+    return null;
+  }
+}
+
+/**
+ * Test-facing shadow wrapper: gate on the shadow flag, then run v3
+ * orchestration for one turn and log the selection set. Production injection
+ * goes through `produce()` → {@link observeTurn} directly (which checks both
+ * flags); this wrapper exists so tests can drive shadow observation in
+ * isolation. Never throws — all failures are logged and swallowed so the live
+ * turn is unaffected.
+ */
+export async function runShadowObservation(
+  conversationId: string,
+  turnIndex: number,
+): Promise<void> {
+  if (!isAssistantFeatureFlagEnabled(MEMORY_V3_SHADOW, getConfig())) return;
+  await observeTurn(conversationId, turnIndex);
+}
+
+/**
+ * The v3 injector. Reads both flags:
+ *   - `memory-v3-live` on → orchestrate, log, render the working-set selection
+ *     into a `<memory>` block, and return it at v2's dynamic-memory placement.
+ *   - `memory-v3-shadow` on (live off) → orchestrate + log only, return `null`.
+ *   - both off → return `null` (no orchestration).
+ *
+ * Empty selection and any failure return `null` (no v3 injection). v2
+ * suppression keys off BOTH the flag AND this return value, so a `null` here
+ * (failure or empty selection) falls back to v2 memory rather than dropping all
+ * memory.
+ */
+const memoryV3Injector: Injector = {
+  name: "memory-v3-shadow",
+  // High order so it sorts last; the live `<memory>` block uses the
+  // after-memory-prefix placement so it lands at the memory boundary regardless
+  // of this sort key, which only orders content-producing injectors.
+  order: 1000,
+  async produce(ctx: PluginTurnContext): Promise<InjectionBlock | null> {
+    const config = getConfig();
+    const live = isAssistantFeatureFlagEnabled(MEMORY_V3_LIVE, config);
+    const shadow = isAssistantFeatureFlagEnabled(MEMORY_V3_SHADOW, config);
+    if (!live && !shadow) return null;
+
+    const result = await observeTurn(ctx.conversationId, ctx.turnIndex);
+    if (!live || !result) return null;
+
+    try {
+      // `renderMemoryBlock` returns "" for an empty selection; inject nothing.
+      const text = await renderMemoryBlock(
+        result.finalInjection,
+        renderV3PageContent,
+      );
+      if (text.length === 0) return null;
+      return {
+        id: MEMORY_V3_BLOCK_ID,
+        text,
+        // Mirror v2's dynamic `<memory>` block placement.
+        placement: "after-memory-prefix",
+      };
+    } catch (err) {
+      log.warn(
+        {
+          err: err instanceof Error ? err.message : String(err),
+          conversationId: ctx.conversationId,
+        },
+        "memory-v3 live render failed (non-fatal) — falling back to v2",
+      );
+      return null;
+    }
+  },
+};
+
+export const memoryV3ShadowPlugin: Plugin = {
+  manifest: {
+    name: "memory-v3-shadow",
+    version: "0.0.1",
+  },
+  injectors: [memoryV3Injector],
+};

package/src/memory/v3/snapshot.ts

@@ -0,0 +1,209 @@
+import { promises as fs } from "node:fs";
+import path from "node:path";
+
+/**
+ * Snapshot / restore for the memory-v3 data directory.
+ *
+ * The v3 data dir (`<workspace>/memory/v3/data/`) holds the tree's durable
+ * state: `leaves/**\/*.md`, `assignments.json`, and `core.json` (see
+ * `resolveDataDir` in `tree.ts`). A future tree-gardening reconciler mutates
+ * these in place; this module lets it take a snapshot first and roll back if a
+ * reconcile pass produces a bad tree.
+ *
+ * Page `leaves:` frontmatter lives outside the data dir (in the v2 page store
+ * at `<workspace>/memory/concepts/`). The reconciler can capture the prior
+ * page->leaves mapping as `pageRefs` and hand it to {@link snapshotDataDir};
+ * it is persisted alongside the data snapshot and returned by
+ * {@link restoreDataDir} so the caller can revert those external frontmatter
+ * edits too.
+ */
+
+/** Files (relative to the data dir) captured by a snapshot. */
+const SNAPSHOT_FILES = ["assignments.json", "core.json"] as const;
+/** Directories (relative to the data dir) captured recursively. */
+const SNAPSHOT_DIRS = ["leaves"] as const;
+/** Filename for the persisted page->leaves diff inside a snapshot. */
+const PAGE_REFS_FILE = "page-refs.json";
+/** Sibling directory (relative to the data dir) that holds snapshots. */
+const SNAPSHOTS_DIRNAME = "v3-snapshots";
+/** Maximum number of snapshots to retain; older ones are pruned. */
+const MAX_RETAINED_SNAPSHOTS = 5;
+
+function snapshotsRoot(dataDir: string): string {
+  return path.join(path.dirname(dataDir), SNAPSHOTS_DIRNAME);
+}
+
+async function copyFileIfExists(src: string, dest: string): Promise<void> {
+  try {
+    await fs.mkdir(path.dirname(dest), { recursive: true });
+    await fs.copyFile(src, dest);
+  } catch (err) {
+    if ((err as NodeJS.ErrnoException).code === "ENOENT") return;
+    throw err;
+  }
+}
+
+/** Copy `src` into `dest` recursively. Returns false if `src` doesn't exist. */
+async function copyDirIfExists(src: string, dest: string): Promise<boolean> {
+  let entries: import("node:fs").Dirent[];
+  try {
+    entries = await fs.readdir(src, { withFileTypes: true });
+  } catch (err) {
+    if ((err as NodeJS.ErrnoException).code === "ENOENT") return false;
+    throw err;
+  }
+  await fs.mkdir(dest, { recursive: true });
+  for (const entry of entries) {
+    const from = path.join(src, entry.name);
+    const to = path.join(dest, entry.name);
+    if (entry.isDirectory()) {
+      await copyDirIfExists(from, to);
+    } else if (entry.isFile()) {
+      await fs.copyFile(from, to);
+    }
+  }
+  return true;
+}
+
+async function removeIfExists(target: string): Promise<void> {
+  await fs.rm(target, { recursive: true, force: true });
+}
+
+function serializePageRefs(pageRefs: Map<string, string[]>): string {
+  return JSON.stringify(Object.fromEntries(pageRefs), null, 2);
+}
+
+function deserializePageRefs(raw: string): Map<string, string[]> {
+  const obj = JSON.parse(raw) as Record<string, string[]>;
+  return new Map(Object.entries(obj));
+}
+
+/**
+ * Copy the v3 data dir's durable state into a timestamped sibling snapshot
+ * directory and return its path.
+ *
+ * @param dataDir Absolute path to the v3 data dir.
+ * @param opts.pageRefs Optional slug->prior-leaves map describing page
+ *   frontmatter that lives outside the data dir; persisted so a later restore
+ *   can revert it.
+ * @param opts.label Snapshot directory name. Defaults to `Date.now()`; pass an
+ *   explicit label to keep tests deterministic.
+ * @returns Absolute path to the created snapshot directory.
+ */
+export async function snapshotDataDir(
+  dataDir: string,
+  opts?: { pageRefs?: Map<string, string[]>; label?: string },
+): Promise<string> {
+  const label = opts?.label ?? String(Date.now());
+  const snapshotPath = path.join(snapshotsRoot(dataDir), label);
+  await fs.mkdir(snapshotPath, { recursive: true });
+
+  for (const file of SNAPSHOT_FILES) {
+    await copyFileIfExists(
+      path.join(dataDir, file),
+      path.join(snapshotPath, file),
+    );
+  }
+  for (const dir of SNAPSHOT_DIRS) {
+    await copyDirIfExists(
+      path.join(dataDir, dir),
+      path.join(snapshotPath, dir),
+    );
+  }
+  if (opts?.pageRefs) {
+    await fs.writeFile(
+      path.join(snapshotPath, PAGE_REFS_FILE),
+      serializePageRefs(opts.pageRefs),
+    );
+  }
+
+  await pruneSnapshots(dataDir);
+  return snapshotPath;
+}
+
+/**
+ * Atomically replace the v3 data dir's durable state from a snapshot.
+ *
+ * Each captured file/dir is staged next to its destination and renamed into
+ * place, so a crash mid-restore never leaves a half-copied tree.
+ *
+ * @param snapshotPath Absolute path to a directory created by
+ *   {@link snapshotDataDir}.
+ * @param dataDir Absolute path to the v3 data dir to overwrite.
+ * @returns The persisted `pageRefs` map (if the snapshot captured one) so the
+ *   caller can revert external page frontmatter.
+ */
+export async function restoreDataDir(
+  snapshotPath: string,
+  dataDir: string,
+): Promise<{ pageRefs?: Map<string, string[]> }> {
+  await fs.mkdir(dataDir, { recursive: true });
+
+  for (const file of SNAPSHOT_FILES) {
+    const src = path.join(snapshotPath, file);
+    const dest = path.join(dataDir, file);
+    const tmp = `${dest}.restore-tmp`;
+    try {
+      await fs.copyFile(src, tmp);
+    } catch (err) {
+      if ((err as NodeJS.ErrnoException).code === "ENOENT") {
+        // Snapshot didn't capture this file; drop any stale destination copy.
+        await removeIfExists(dest);
+        continue;
+      }
+      throw err;
+    }
+    await fs.rename(tmp, dest);
+  }
+
+  for (const dir of SNAPSHOT_DIRS) {
+    const src = path.join(snapshotPath, dir);
+    const dest = path.join(dataDir, dir);
+    const tmp = `${dest}.restore-tmp`;
+    await removeIfExists(tmp);
+    const captured = await copyDirIfExists(src, tmp);
+    if (!captured) {
+      // Snapshot didn't capture this dir; drop any stale destination copy.
+      await removeIfExists(dest);
+      continue;
+    }
+    await removeIfExists(dest);
+    await fs.rename(tmp, dest);
+  }
+
+  const pageRefsPath = path.join(snapshotPath, PAGE_REFS_FILE);
+  try {
+    const raw = await fs.readFile(pageRefsPath, "utf8");
+    return { pageRefs: deserializePageRefs(raw) };
+  } catch (err) {
+    if ((err as NodeJS.ErrnoException).code === "ENOENT") return {};
+    throw err;
+  }
+}
+
+/**
+ * Delete all but the most recent {@link MAX_RETAINED_SNAPSHOTS} snapshot
+ * directories. Snapshots sort lexicographically by their label; numeric
+ * timestamps and zero-padded labels both order chronologically.
+ */
+async function pruneSnapshots(dataDir: string): Promise<void> {
+  const root = snapshotsRoot(dataDir);
+  let entries: import("node:fs").Dirent[];
+  try {
+    entries = await fs.readdir(root, { withFileTypes: true });
+  } catch (err) {
+    if ((err as NodeJS.ErrnoException).code === "ENOENT") return;
+    throw err;
+  }
+  const dirs = entries
+    .filter((e) => e.isDirectory())
+    .map((e) => e.name)
+    .sort();
+  const stale = dirs.slice(
+    0,
+    Math.max(0, dirs.length - MAX_RETAINED_SNAPSHOTS),
+  );
+  for (const name of stale) {
+    await removeIfExists(path.join(root, name));
+  }
+}