@vellumai/assistant 0.8.6 → 0.8.7

package/src/memory/v3/__tests__/selection-log-store.test.ts

@@ -0,0 +1,179 @@
+/**
+ * Tests for `assistant/src/memory/v3/selection-log-store.ts`.
+ *
+ * Asserts the inspector read path over `memory_v3_selections`:
+ *   - the selection for the exact turn;
+ *   - null for a null turn, and null for a turn with no v3 rows;
+ *   - NO fallback to another turn (would misattribute a future turn's
+ *     selection to the inspected message);
+ *   - source/pinned mapping and the rendered `<memory>` block;
+ *   - `live` / `shadow` reflect the flag resolver.
+ *
+ * `mock.module` is process-global and leaks into sibling files in a
+ * `bun test src/memory/v3/` run, so every stub DELEGATES to the real
+ * implementation unless this test is actively running (`storeMockActive`,
+ * toggled in beforeEach/afterAll). Mirrors `shadow-plugin.test.ts`.
+ */
+
+import { Database } from "bun:sqlite";
+import { afterAll, beforeEach, describe, expect, mock, test } from "bun:test";
+
+import { drizzle } from "drizzle-orm/bun-sqlite";
+
+import { migrateAddMemoryV3Selections } from "../../migrations/268-add-memory-v3-selections.js";
+import * as schema from "../../schema.js";
+
+const realFlags = {
+  ...(await import("../../../config/assistant-feature-flags.js")),
+};
+const realLoader = { ...(await import("../../../config/loader.js")) };
+const realDb = { ...(await import("../../db-connection.js")) };
+const realPageContent = { ...(await import("../page-content.js")) };
+
+let storeMockActive = false;
+let liveEnabled = false;
+let shadowEnabled = false;
+
+let testSqlite: Database;
+let testDb = makeDb();
+function makeDb() {
+  testSqlite = new Database(":memory:");
+  testSqlite.exec("PRAGMA journal_mode=WAL");
+  const db = drizzle(testSqlite, { schema });
+  migrateAddMemoryV3Selections(db);
+  return db;
+}
+
+function seed(
+  conversationId: string,
+  turn: number,
+  rows: Array<{ slug: string; source: string; pinned?: boolean }>,
+): void {
+  const stmt = testSqlite.query(
+    `INSERT INTO memory_v3_selections (conversation_id, turn, slug, source, pinned, created_at) VALUES (?, ?, ?, ?, ?, ?)`,
+  );
+  for (const r of rows) {
+    stmt.run(
+      conversationId,
+      turn,
+      r.slug,
+      r.source,
+      r.pinned ? 1 : 0,
+      1000 + turn,
+    );
+  }
+}
+
+mock.module("../../../config/assistant-feature-flags.js", () => ({
+  ...realFlags,
+  isAssistantFeatureFlagEnabled: (key: string, config: unknown) =>
+    storeMockActive
+      ? key === "memory-v3-live"
+        ? liveEnabled
+        : key === "memory-v3-shadow"
+          ? shadowEnabled
+          : false
+      : realFlags.isAssistantFeatureFlagEnabled(
+          key as Parameters<typeof realFlags.isAssistantFeatureFlagEnabled>[0],
+          config as Parameters<
+            typeof realFlags.isAssistantFeatureFlagEnabled
+          >[1],
+        ),
+}));
+
+mock.module("../../../config/loader.js", () => ({
+  ...realLoader,
+  getConfig: () => (storeMockActive ? {} : realLoader.getConfig()),
+}));
+
+mock.module("../../db-connection.js", () => ({
+  ...realDb,
+  getDb: () => (storeMockActive ? testDb : realDb.getDb()),
+  getSqliteFrom: (db: unknown) =>
+    storeMockActive
+      ? testSqlite
+      : realDb.getSqliteFrom(db as Parameters<typeof realDb.getSqliteFrom>[0]),
+}));
+
+mock.module("../page-content.js", () => ({
+  ...realPageContent,
+  renderV3PageContent: async (slug: string) =>
+    storeMockActive
+      ? `body for ${slug}`
+      : realPageContent.renderV3PageContent(slug),
+}));
+
+const { getMemoryV3SelectionForInspector } =
+  await import("../selection-log-store.js");
+
+beforeEach(() => {
+  storeMockActive = true;
+  liveEnabled = false;
+  shadowEnabled = false;
+  testDb = makeDb();
+});
+
+afterAll(() => {
+  storeMockActive = false;
+});
+
+describe("getMemoryV3SelectionForInspector", () => {
+  test("returns null for a null/undefined turn", async () => {
+    seed("conv-x", 3, [{ slug: "domain-a/page-1", source: "l1+l2" }]);
+    expect(await getMemoryV3SelectionForInspector("conv-x", null)).toBeNull();
+    expect(
+      await getMemoryV3SelectionForInspector("conv-x", undefined),
+    ).toBeNull();
+  });
+
+  test("returns the selection for the exact turn", async () => {
+    seed("conv-2", 2, [{ slug: "domain-a/page-1", source: "l1+l2" }]);
+    seed("conv-2", 7, [{ slug: "domain-b/page-9", source: "l1+l2" }]);
+
+    const log = await getMemoryV3SelectionForInspector("conv-2", 2);
+    expect(log?.turn).toBe(2);
+    expect(log?.selections.map((s) => s.slug)).toEqual(["domain-a/page-1"]);
+  });
+
+  test("returns null when the turn has no v3 rows", async () => {
+    seed("conv-1", 3, [{ slug: "domain-a/page-1", source: "l1+l2" }]);
+    expect(await getMemoryV3SelectionForInspector("conv-1", 4)).toBeNull();
+  });
+
+  test("does NOT fall back to another turn for an unmatched lookup", async () => {
+    // Turn 5 has rows, but inspecting turn 3 (no rows) must return null —
+    // never turn 5's selection, which would misattribute it to turn 3.
+    seed("conv-3", 5, [{ slug: "domain-a/page-1", source: "l1+l2" }]);
+    expect(await getMemoryV3SelectionForInspector("conv-3", 3)).toBeNull();
+  });
+
+  test("maps source/pinned and renders the <memory> block", async () => {
+    seed("conv-4", 1, [
+      { slug: "domain-a/page-1", source: "core+l2", pinned: true },
+      { slug: "domain-b/page-2", source: "carry-forward", pinned: false },
+    ]);
+
+    const log = await getMemoryV3SelectionForInspector("conv-4", 1);
+    expect(log?.selections).toEqual([
+      { slug: "domain-a/page-1", source: "core+l2", pinned: true },
+      { slug: "domain-b/page-2", source: "carry-forward", pinned: false },
+    ]);
+    expect(log?.injectedText).toContain("<memory>");
+    expect(log?.injectedText).toContain("body for domain-a/page-1");
+    expect(log?.injectedText).toContain("body for domain-b/page-2");
+  });
+
+  test("live/shadow reflect the flag resolver", async () => {
+    seed("conv-5", 1, [{ slug: "domain-a/page-1", source: "l1+l2" }]);
+
+    const off = await getMemoryV3SelectionForInspector("conv-5", 1);
+    expect(off?.live).toBe(false);
+    expect(off?.shadow).toBe(false);
+
+    liveEnabled = true;
+    shadowEnabled = true;
+    const on = await getMemoryV3SelectionForInspector("conv-5", 1);
+    expect(on?.live).toBe(true);
+    expect(on?.shadow).toBe(true);
+  });
+});

package/src/memory/v3/__tests__/selector.test.ts

@@ -0,0 +1,404 @@
+/**
+ * Tests for `assistant/src/memory/v3/selector.ts`.
+ *
+ * Coverage matrix:
+ *   - Returned IDs map to the right member slugs by 1-based index, with
+ *     `pinned` driven by `pinned_ids`.
+ *   - Omitted `ids` → ALL members of the leaf (recall-safe).
+ *   - Explicit `ids: []` → no pages (deliberate abstention).
+ *   - No provider / missing tool_use / schema mismatch / throw → ALL members.
+ *   - The per-leaf `<pages>` prefix is byte-identical across two calls with
+ *     different turns (the cache invariant).
+ *   - `selectAcrossLeaves` flattens per-leaf results and never exceeds the
+ *     configured concurrency.
+ *
+ * The provider is stubbed so no network calls fire.
+ */
+
+import { beforeEach, describe, expect, mock, test } from "bun:test";
+
+import type {
+  Message,
+  Provider,
+  ProviderResponse,
+  SendMessageOptions,
+  ToolUseContent,
+} from "../../../providers/types.js";
+import type {
+  LeafNode,
+  LeafPath,
+  LeafTree,
+  Slug,
+  TurnContext,
+} from "../types.js";
+
+// ---------------------------------------------------------------------------
+// Mocks installed BEFORE the selector import so the module observes them at
+// load time.
+// ---------------------------------------------------------------------------
+
+let providerStub: Provider | null = null;
+
+interface ProviderCall {
+  messages: Message[];
+  options: SendMessageOptions | undefined;
+}
+const providerCalls: ProviderCall[] = [];
+
+mock.module("../../../providers/provider-send-message.js", () => ({
+  getConfiguredProvider: async () => providerStub,
+  extractToolUse: (response: ProviderResponse) =>
+    response.content.find((b): b is ToolUseContent => b.type === "tool_use"),
+}));
+
+mock.module("../../../util/logger.js", () => ({
+  getLogger: () =>
+    new Proxy({} as Record<string, unknown>, {
+      get: (_t, prop) => (prop === "child" ? () => ({}) : () => {}),
+    }),
+}));
+
+const { selectFromLeaf, selectAcrossLeaves } = await import("../selector.js");
+
+// ---------------------------------------------------------------------------
+// Helpers.
+// ---------------------------------------------------------------------------
+
+function makeProvider(response: ProviderResponse): Provider {
+  return {
+    name: "stub",
+    sendMessage: async (messages, options) => {
+      providerCalls.push({ messages, options });
+      return response;
+    },
+  };
+}
+
+function toolUseResponse(input: Record<string, unknown>): ProviderResponse {
+  return {
+    model: "stub-model",
+    stopReason: "tool_use",
+    usage: { inputTokens: 0, outputTokens: 0 },
+    content: [{ type: "tool_use", id: "tu-1", name: "select_pages", input }],
+  };
+}
+
+function makeLeaf(path: LeafPath, members: Slug[]): LeafNode {
+  return {
+    path,
+    frontmatter: { path, in_core: false },
+    description: `description for ${path}`,
+    members,
+    domain: path.split("/")[0],
+  };
+}
+
+/** Tree with two leaves; `people/alice` has three member pages. */
+function makeTree(): LeafTree {
+  const leaves: LeafNode[] = [
+    makeLeaf("people/alice", ["alice-bio", "alice-1on1", "alice-feedback"]),
+    makeLeaf("projects/atlas", ["atlas-roadmap", "atlas-status"]),
+  ];
+  const byPage = new Map<Slug, LeafPath[]>();
+  for (const leaf of leaves) {
+    for (const slug of leaf.members) {
+      byPage.set(slug, [...(byPage.get(slug) ?? []), leaf.path]);
+    }
+  }
+  return {
+    leaves: new Map(leaves.map((n) => [n.path, n])),
+    byPage,
+  };
+}
+
+const ALICE_MEMBERS = ["alice-bio", "alice-1on1", "alice-feedback"];
+
+function makeTurn(currentMessage: string): TurnContext {
+  return {
+    conversationId: "conv-xyz",
+    turnNumber: 1,
+    currentMessage,
+    recentContext: "earlier we talked about the timeline",
+  };
+}
+
+/** Deterministic summary stub so the rendered prefix is reproducible. */
+const summaryOf = async (slug: Slug): Promise<string> => `summary of ${slug}`;
+
+beforeEach(() => {
+  providerStub = null;
+  providerCalls.length = 0;
+});
+
+// ---------------------------------------------------------------------------
+// selectFromLeaf — id mapping.
+// ---------------------------------------------------------------------------
+
+describe("selectFromLeaf — id mapping", () => {
+  test("returned IDs map to member slugs, pinned driven by pinned_ids", async () => {
+    providerStub = makeProvider(
+      toolUseResponse({ ids: [3, 1], pinned_ids: [1] }),
+    );
+    const result = await selectFromLeaf(
+      "people/alice",
+      makeTurn("how's alice?"),
+      makeTree(),
+      summaryOf,
+    );
+    expect(result).toEqual([
+      { slug: "alice-feedback", pinned: false },
+      { slug: "alice-bio", pinned: true },
+    ]);
+  });
+
+  test("omitted ids selects ALL members (recall-safe)", async () => {
+    providerStub = makeProvider(toolUseResponse({}));
+    const result = await selectFromLeaf(
+      "people/alice",
+      makeTurn("anything"),
+      makeTree(),
+      summaryOf,
+    );
+    expect(result).toEqual(
+      ALICE_MEMBERS.map((slug) => ({ slug, pinned: false })),
+    );
+  });
+
+  test("explicit empty ids selects no pages (abstention)", async () => {
+    providerStub = makeProvider(toolUseResponse({ ids: [] }));
+    const result = await selectFromLeaf(
+      "people/alice",
+      makeTurn("nothing relevant"),
+      makeTree(),
+      summaryOf,
+    );
+    expect(result).toEqual([]);
+  });
+
+  test("out-of-range and duplicate IDs are ignored without throwing", async () => {
+    providerStub = makeProvider(toolUseResponse({ ids: [2, 99, 0, -1, 2] }));
+    const result = await selectFromLeaf(
+      "people/alice",
+      makeTurn("the 1:1"),
+      makeTree(),
+      summaryOf,
+    );
+    expect(result).toEqual([{ slug: "alice-1on1", pinned: false }]);
+  });
+
+  test("empty leaf returns no pages and never calls the provider", async () => {
+    providerStub = makeProvider(toolUseResponse({ ids: [1] }));
+    const result = await selectFromLeaf(
+      "people/empty",
+      makeTurn("hi"),
+      {
+        leaves: new Map([["people/empty", makeLeaf("people/empty", [])]]),
+        byPage: new Map(),
+      },
+      summaryOf,
+    );
+    expect(result).toEqual([]);
+    expect(providerCalls).toHaveLength(0);
+  });
+});
+
+// ---------------------------------------------------------------------------
+// selectFromLeaf — recall-safe fallbacks.
+// ---------------------------------------------------------------------------
+
+describe("selectFromLeaf — recall-safe fallbacks", () => {
+  const allAlice = ALICE_MEMBERS.map((slug) => ({ slug, pinned: false }));
+
+  test("no provider → ALL members", async () => {
+    providerStub = null;
+    const result = await selectFromLeaf(
+      "people/alice",
+      makeTurn("x"),
+      makeTree(),
+      summaryOf,
+    );
+    expect(result).toEqual(allAlice);
+  });
+
+  test("missing tool_use → ALL members", async () => {
+    providerStub = makeProvider({
+      model: "stub-model",
+      stopReason: "end_turn",
+      usage: { inputTokens: 0, outputTokens: 0 },
+      content: [{ type: "text", text: "no tool call" }],
+    });
+    const result = await selectFromLeaf(
+      "people/alice",
+      makeTurn("x"),
+      makeTree(),
+      summaryOf,
+    );
+    expect(result).toEqual(allAlice);
+  });
+
+  test("schema mismatch → ALL members", async () => {
+    providerStub = makeProvider(toolUseResponse({ ids: "not-an-array" }));
+    const result = await selectFromLeaf(
+      "people/alice",
+      makeTurn("x"),
+      makeTree(),
+      summaryOf,
+    );
+    expect(result).toEqual(allAlice);
+  });
+
+  test("provider throw → ALL members", async () => {
+    providerStub = {
+      name: "throwing",
+      sendMessage: async () => {
+        throw new Error("boom");
+      },
+    };
+    const result = await selectFromLeaf(
+      "people/alice",
+      makeTurn("x"),
+      makeTree(),
+      summaryOf,
+    );
+    expect(result).toEqual(allAlice);
+  });
+});
+
+// ---------------------------------------------------------------------------
+// selectFromLeaf — request shape and cache invariant.
+// ---------------------------------------------------------------------------
+
+describe("selectFromLeaf — request shape", () => {
+  test("forces tool_choice to select_pages with the v3 L2 call site", async () => {
+    providerStub = makeProvider(toolUseResponse({ ids: [1] }));
+    await selectFromLeaf(
+      "people/alice",
+      makeTurn("alice?"),
+      makeTree(),
+      summaryOf,
+    );
+
+    expect(providerCalls).toHaveLength(1);
+    const [call] = providerCalls;
+    const cfg = call.options?.config as Record<string, unknown>;
+    expect(cfg?.callSite).toBe("memoryV3SelectL2");
+    expect(cfg?.tool_choice).toEqual({ type: "tool", name: "select_pages" });
+    expect(call.options?.tools?.[0]?.name).toBe("select_pages");
+  });
+
+  test("pages block is the first content block with an ephemeral cache breakpoint", async () => {
+    providerStub = makeProvider(toolUseResponse({ ids: [1] }));
+    await selectFromLeaf(
+      "people/alice",
+      makeTurn("alice?"),
+      makeTree(),
+      summaryOf,
+    );
+
+    const [blockA, blockB] = providerCalls[0].messages[0].content as Array<{
+      type: string;
+      text: string;
+      cache_control?: { type: string };
+    }>;
+    expect(blockA.type).toBe("text");
+    expect(blockA.text).toContain("<leaf>people/alice</leaf>");
+    expect(blockA.text).toContain("<pages>");
+    expect(blockA.text).toContain("[1] alice-bio — summary of alice-bio");
+    expect(blockA.cache_control).toEqual({ type: "ephemeral" });
+
+    expect(blockB.type).toBe("text");
+    expect(blockB.text).toContain("<current_message>alice?</current_message>");
+    expect(blockB.text).toContain("<recent_context>");
+    expect(blockB.cache_control).toBeUndefined();
+  });
+
+  test("system prompt mentions pinned (locks the pinning commitment)", async () => {
+    providerStub = makeProvider(toolUseResponse({ ids: [1] }));
+    await selectFromLeaf("people/alice", makeTurn("x"), makeTree(), summaryOf);
+    expect(providerCalls[0].options?.systemPrompt).toMatch(/pinned/);
+  });
+});
+
+describe("selectFromLeaf — per-leaf cache invariant", () => {
+  test("the pages prefix is byte-identical across two calls with different turns", async () => {
+    const tree = makeTree();
+    providerStub = makeProvider(toolUseResponse({ ids: [1] }));
+
+    await selectFromLeaf(
+      "people/alice",
+      makeTurn("first query"),
+      tree,
+      summaryOf,
+    );
+    await selectFromLeaf(
+      "people/alice",
+      makeTurn("totally different second query"),
+      tree,
+      summaryOf,
+    );
+
+    const prefixA = (
+      providerCalls[0].messages[0].content[0] as { text: string }
+    ).text;
+    const prefixB = (
+      providerCalls[1].messages[0].content[0] as { text: string }
+    ).text;
+    expect(prefixA).toBe(prefixB);
+  });
+});
+
+// ---------------------------------------------------------------------------
+// selectAcrossLeaves — fan-out and bounded concurrency.
+// ---------------------------------------------------------------------------
+
+describe("selectAcrossLeaves", () => {
+  test("flattens per-leaf selections", async () => {
+    // The selector forces the tool, so omitted ids → all members per leaf.
+    providerStub = makeProvider(toolUseResponse({}));
+    const result = await selectAcrossLeaves(
+      ["people/alice", "projects/atlas"],
+      makeTurn("x"),
+      makeTree(),
+      summaryOf,
+    );
+    expect(result.map((p) => p.slug)).toEqual([
+      "alice-bio",
+      "alice-1on1",
+      "alice-feedback",
+      "atlas-roadmap",
+      "atlas-status",
+    ]);
+  });
+
+  test("never exceeds the configured concurrency", async () => {
+    let inFlight = 0;
+    let maxInFlight = 0;
+    let calls = 0;
+    providerStub = {
+      name: "tracking",
+      sendMessage: async () => {
+        calls++;
+        inFlight++;
+        maxInFlight = Math.max(maxInFlight, inFlight);
+        await new Promise((r) => setTimeout(r, 5));
+        inFlight--;
+        return toolUseResponse({ ids: [] });
+      },
+    };
+
+    // Build a tree with eight single-member leaves so there are eight calls.
+    const leaves: LeafNode[] = Array.from({ length: 8 }, (_, i) =>
+      makeLeaf(`misc/leaf-${i}`, [`page-${i}`]),
+    );
+    const tree: LeafTree = {
+      leaves: new Map(leaves.map((n) => [n.path, n])),
+      byPage: new Map(),
+    };
+    const paths = leaves.map((l) => l.path);
+
+    await selectAcrossLeaves(paths, makeTurn("x"), tree, summaryOf, 3);
+    expect(maxInFlight).toBeLessThanOrEqual(3);
+    expect(maxInFlight).toBeGreaterThan(1);
+    expect(calls).toBe(8);
+  });
+});