@vellumai/assistant 0.8.4 → 0.8.6

package/src/memory/v3/__tests__/shadow-diff.test.ts

@@ -0,0 +1,225 @@
+import { describe, expect, test } from "bun:test";
+
+import type { MemoryV2ConceptRowRecord } from "../../memory-v2-activation-log-store.js";
+import { computeShadowDiff, type ShadowDiffLogRow } from "../shadow-diff.js";
+
+const ZERO = {
+  finalActivation: 0,
+  ownActivation: 0,
+  priorActivation: 0,
+  simUser: 0,
+  simAssistant: 0,
+  simNow: 0,
+  simUserRerankBoost: 0,
+  simAssistantRerankBoost: 0,
+  inRerankPool: false,
+  spreadContribution: 0,
+} as const;
+
+/** A v2 router concept row with the given selection status. */
+function v2concept(
+  slug: string,
+  status: MemoryV2ConceptRowRecord["status"],
+): MemoryV2ConceptRowRecord {
+  return { slug, status, source: "tier3:0", ...ZERO };
+}
+
+/** A v3 shadow concept row (always status 'injected') tagged with its lane. */
+function v3concept(slug: string, lane: string): MemoryV2ConceptRowRecord {
+  return { slug, status: "injected", source: "router", lane, ...ZERO };
+}
+
+function shadowRow(
+  conversationId: string,
+  createdAt: number,
+  slugLanes: Array<[string, string]>,
+): ShadowDiffLogRow {
+  return {
+    conversationId,
+    createdAt,
+    concepts: slugLanes.map(([slug, lane]) => v3concept(slug, lane)),
+  };
+}
+
+function routerRow(
+  conversationId: string,
+  createdAt: number,
+  opts: { injected?: string[]; cached?: string[]; rejected?: string[] },
+): ShadowDiffLogRow {
+  return {
+    conversationId,
+    createdAt,
+    concepts: [
+      ...(opts.injected ?? []).map((s) => v2concept(s, "injected")),
+      ...(opts.cached ?? []).map((s) => v2concept(s, "in_context")),
+      ...(opts.rejected ?? []).map((s) => v2concept(s, "not_injected")),
+    ],
+  };
+}
+
+const OPTS = { toleranceMs: 10_000, detailLimit: 50 };
+
+describe("computeShadowDiff", () => {
+  test("pairs nearest router row and diffs fresh-injected vs shadow sets", () => {
+    const shadow = [
+      shadowRow("conv1", 1000, [
+        ["a", "tree"],
+        ["b", "edge"],
+        ["c", "hot"],
+      ]),
+    ];
+    const router = [
+      routerRow("conv1", 1500, {
+        injected: ["a", "b", "x", "y"],
+        cached: ["cached1", "cached2"],
+        rejected: ["rej1"],
+      }),
+    ];
+
+    const result = computeShadowDiff(shadow, router, OPTS);
+
+    expect(result.turnsCompared).toBe(1);
+    expect(result.unpaired).toHaveLength(0);
+
+    const turn = result.turns[0]!;
+    expect(turn.deltaMs).toBe(500);
+    expect(turn.v2Count).toBe(4);
+    expect(turn.v3Count).toBe(3);
+    expect(turn.v2CachedCount).toBe(2);
+    expect(turn.overlap).toEqual(["a", "b"]);
+    expect(turn.v3Only).toEqual(["c"]);
+    expect(turn.v2Only).toEqual(["x", "y"]);
+    // cached / rejected v2 rows are NOT counted as dropped picks.
+    expect(turn.v2Only).not.toContain("cached1");
+    expect(turn.v2Only).not.toContain("rej1");
+    expect(turn.jaccard).toBeCloseTo(2 / 5, 10);
+  });
+
+  test("attributes overlap and v3-only to the v3 provenance lane", () => {
+    const shadow = [
+      shadowRow("conv1", 1000, [
+        ["a", "tree"],
+        ["b", "edge"],
+        ["c", "hot"],
+      ]),
+    ];
+    const router = [routerRow("conv1", 1100, { injected: ["a", "b"] })];
+
+    const { agg } = computeShadowDiff(shadow, router, OPTS);
+
+    expect(agg.overlapByLane).toEqual({ tree: 1, edge: 1 });
+    expect(agg.v3OnlyByLane).toEqual({ hot: 1 });
+  });
+
+  test("defaults a missing lane to 'unknown'", () => {
+    const shadow: ShadowDiffLogRow[] = [
+      {
+        conversationId: "conv1",
+        createdAt: 1000,
+        concepts: [
+          { slug: "a", status: "injected", source: "router", ...ZERO },
+        ],
+      },
+    ];
+    const router = [routerRow("conv1", 1100, { injected: [] })];
+
+    const { agg } = computeShadowDiff(shadow, router, OPTS);
+
+    expect(agg.v3OnlyByLane).toEqual({ unknown: 1 });
+  });
+
+  test("sends a shadow row with no router row within tolerance to unpaired", () => {
+    const shadow = [shadowRow("conv1", 1000, [["a", "tree"]])];
+    const router = [routerRow("conv1", 1000 + 20_000, { injected: ["a"] })];
+
+    const result = computeShadowDiff(shadow, router, OPTS);
+
+    expect(result.turnsCompared).toBe(0);
+    expect(result.unpaired).toEqual([
+      { conversationId: "conv1", shadowAt: 1000, v3Count: 1 },
+    ]);
+  });
+
+  test("greedily pairs each router row to at most one shadow row", () => {
+    const shadow = [
+      shadowRow("conv1", 1000, [["a", "tree"]]),
+      shadowRow("conv1", 5000, [["b", "edge"]]),
+    ];
+    const router = [
+      routerRow("conv1", 1200, { injected: ["a", "z"] }),
+      routerRow("conv1", 5200, { injected: ["b"] }),
+    ];
+
+    const result = computeShadowDiff(shadow, router, OPTS);
+
+    expect(result.turnsCompared).toBe(2);
+    const byShadow = new Map(result.turns.map((t) => [t.shadowAt, t]));
+    expect(byShadow.get(1000)!.routerAt).toBe(1200);
+    expect(byShadow.get(1000)!.v2Only).toEqual(["z"]);
+    expect(byShadow.get(5000)!.routerAt).toBe(5200);
+    expect(byShadow.get(5000)!.overlap).toEqual(["b"]);
+  });
+
+  test("caps per-turn detail newest-first but aggregates over all turns", () => {
+    const shadow = [
+      shadowRow("c1", 1000, [["a", "tree"]]),
+      shadowRow("c2", 2000, [["b", "tree"]]),
+      shadowRow("c3", 3000, [["c", "tree"]]),
+    ];
+    const router = [
+      routerRow("c1", 1000, { injected: ["a"] }),
+      routerRow("c2", 2000, { injected: ["b"] }),
+      routerRow("c3", 3000, { injected: ["c"] }),
+    ];
+
+    const result = computeShadowDiff(shadow, router, {
+      toleranceMs: 10_000,
+      detailLimit: 2,
+    });
+
+    expect(result.turnsCompared).toBe(3);
+    expect(result.turns.map((t) => t.shadowAt)).toEqual([3000, 2000]);
+  });
+
+  test("ranks the most-dropped and most-added slugs across turns", () => {
+    const shadow = [
+      shadowRow("c1", 1000, [["extra", "edge"]]),
+      shadowRow("c2", 2000, [["extra", "edge"]]),
+    ];
+    const router = [
+      routerRow("c1", 1000, { injected: ["dropped", "kept"] }),
+      routerRow("c2", 2000, { injected: ["dropped"] }),
+    ];
+
+    const { agg } = computeShadowDiff(shadow, router, OPTS);
+
+    expect(agg.totalV3Only).toBe(2);
+    expect(agg.totalV2Only).toBe(3);
+    expect(agg.v3OnlyTop[0]).toEqual({ slug: "extra", count: 2 });
+    expect(agg.v2OnlyTop[0]).toEqual({ slug: "dropped", count: 2 });
+    expect(agg.meanV3).toBeCloseTo(1, 10);
+    expect(agg.meanV2).toBeCloseTo(1.5, 10);
+  });
+
+  test("returns zeroed aggregates with no rows", () => {
+    const result = computeShadowDiff([], [], OPTS);
+    expect(result.turnsCompared).toBe(0);
+    expect(result.shadowRows).toBe(0);
+    expect(result.agg.meanJaccard).toBe(0);
+    expect(result.agg.v2OnlyTop).toEqual([]);
+    expect(result.turns).toEqual([]);
+  });
+
+  test("treats two empty selections as jaccard 0, not NaN", () => {
+    const shadow: ShadowDiffLogRow[] = [
+      { conversationId: "c1", createdAt: 1000, concepts: [] },
+    ];
+    const router = [routerRow("c1", 1000, { cached: ["only-cached"] })];
+
+    const result = computeShadowDiff(shadow, router, OPTS);
+
+    expect(result.turnsCompared).toBe(1);
+    expect(result.turns[0]!.jaccard).toBe(0);
+    expect(result.turns[0]!.v2CachedCount).toBe(1);
+  });
+});

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

@@ -0,0 +1,398 @@
+/**
+ * Tests for the live-shadow `memoryRetrieval` middleware
+ * (`assistant/src/memory/v3/shadow-middleware.ts`).
+ *
+ * The critical invariant this PR guarantees: with `memory.v3.shadow` off
+ * (the default), the middleware is a byte-for-byte pass-through — it returns
+ * the downstream `MemoryResult` unchanged, never calls the v3 loop, and never
+ * writes a log row. With the flag on, it runs the v3 loop alongside the
+ * default path, logs v3's selection as `mode='v3_shadow'`, and STILL returns
+ * the unchanged downstream result (v2 injected, never v3). A v3 failure is
+ * swallowed and the turn result is unaffected.
+ *
+ * Everything the middleware reaches (config, the v3 loop, the activation-log
+ * store, message/now/everInjected reads) is stubbed via `mock.module` — no
+ * real LLM, no real workspace DB.
+ */
+
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+
+import { makeMockLogger } from "../../../__tests__/helpers/mock-logger.js";
+import type { AssistantConfig } from "../../../config/schema.js";
+import type { TrustContext } from "../../../daemon/trust-context.js";
+import type {
+  MemoryArgs,
+  MemoryResult,
+  TurnContext,
+} from "../../../plugins/types.js";
+import type { RecordMemoryV2ActivationLogParams } from "../../memory-v2-activation-log-store.js";
+import type {
+  RetrievalInput,
+  RetrievalOutput,
+} from "../../v2/harness/retriever.js";
+
+mock.module("../../../util/logger.js", () => ({
+  getLogger: () => makeMockLogger(),
+}));
+
+// Mirror the real `isUntrustedTrustClass` predicate (anything that is not the
+// guardian is untrusted) without pulling the resolver's contact-store graph
+// into the test module set.
+mock.module("../../../runtime/actor-trust-resolver.js", () => ({
+  isUntrustedTrustClass: (trustClass: string | undefined) =>
+    trustClass === "trusted_contact" ||
+    trustClass === "unknown" ||
+    trustClass === undefined,
+}));
+
+// ── Mutable test doubles, rewired per test ───────────────────────────────
+
+/** Drives `config.memory.v3.{enabled,shadow}` and `historical_pairs`. */
+let v3Enabled = false;
+let v3Shadow = false;
+/** When false, omit the `memory.v3` block entirely (mirrors configs built
+ * outside the Zod schema, e.g. agent-loop test fixtures). */
+let v3Present = true;
+
+function makeConfig(): AssistantConfig {
+  return {
+    memory: {
+      v2: { router: { historical_pairs: 1 } },
+      ...(v3Present ? { v3: { enabled: v3Enabled, shadow: v3Shadow } } : {}),
+    },
+  } as unknown as AssistantConfig;
+}
+
+/** Captured `runRetrievalLoop` invocations. */
+const loopCalls: Array<{ input: RetrievalInput }> = [];
+/** Behavior of the stubbed loop — overridden per test. */
+let loopImpl: (
+  input: RetrievalInput,
+) => Promise<RetrievalOutput> = async () => ({
+  selectedSlugs: [],
+  sourceBySlug: new Map(),
+  trace: { passes: [] },
+  cost: { ms: 0 },
+  failureReason: null,
+});
+
+/** Captured `recordMemoryV2ActivationLog` calls. */
+const logCalls: RecordMemoryV2ActivationLogParams[] = [];
+
+mock.module("../../../config/loader.js", () => ({
+  getConfig: () => makeConfig(),
+}));
+mock.module("../../../util/platform.js", () => ({
+  getWorkspaceDir: () => "/tmp/shadow-test-workspace",
+}));
+// Chainable drizzle-query stub: every builder method returns the same object
+// and `.all()` yields the seeded rows. The shadow middleware reads recent
+// messages via `db.select(...).from(...).where(...).orderBy(...).limit(...).all()`.
+const messageRows: Array<{ role: string; content: string }> = [
+  {
+    role: "user",
+    content: JSON.stringify([{ type: "text", text: "hello memory" }]),
+  },
+];
+function makeFakeDb(): never {
+  const builder: Record<string, unknown> = {};
+  for (const m of ["select", "from", "where", "orderBy", "limit"]) {
+    builder[m] = () => builder;
+  }
+  builder.all = () => messageRows.slice();
+  return builder as never;
+}
+mock.module("../../db-connection.js", () => ({
+  getDb: () => makeFakeDb(),
+}));
+mock.module("../../v2/now-text.js", () => ({
+  loadNowText: async () => "NOW context",
+}));
+mock.module("../../v2/activation-store.js", () => ({
+  hydrate: async () => ({ everInjected: [{ slug: "old/page", turn: 0 }] }),
+}));
+mock.module("../loop.js", () => ({
+  runRetrievalLoop: async (input: RetrievalInput): Promise<RetrievalOutput> => {
+    loopCalls.push({ input });
+    return loopImpl(input);
+  },
+}));
+mock.module("../../memory-v2-activation-log-store.js", () => ({
+  recordMemoryV2ActivationLog: (params: RecordMemoryV2ActivationLogParams) => {
+    logCalls.push(params);
+  },
+}));
+
+const { memoryV3ShadowMiddleware } = await import("../shadow-middleware.js");
+
+// ── Fixtures ─────────────────────────────────────────────────────────────
+
+const trust: TrustContext = {
+  sourceChannel: "vellum",
+  trustClass: "guardian",
+};
+
+function makeCtx(): TurnContext {
+  return {
+    requestId: "req-shadow-test",
+    conversationId: "conv-shadow",
+    turnIndex: 3,
+    trust,
+  };
+}
+
+function makeArgs(
+  signal?: AbortSignal,
+  trustContext: TrustContext | undefined = trust,
+): MemoryArgs {
+  return {
+    conversationId: "conv-shadow",
+    trustContext,
+    turnIndex: 3,
+    signal: signal ?? new AbortController().signal,
+  };
+}
+
+/** The unchanged downstream (v2/default) result the terminal returns. */
+const DOWNSTREAM_RESULT: MemoryResult = {
+  pkbContent: "pkb",
+  nowContent: "now",
+  memoryGraphBlocks: [{ kind: "default.graph" }],
+};
+
+/** Flush the detached shadow chain (microtasks + a macrotask hop). */
+async function flush(): Promise<void> {
+  await new Promise((resolve) => setTimeout(resolve, 0));
+  await Promise.resolve();
+}
+
+beforeEach(() => {
+  v3Enabled = false;
+  v3Shadow = false;
+  v3Present = true;
+  loopCalls.length = 0;
+  logCalls.length = 0;
+  loopImpl = async () => ({
+    selectedSlugs: [],
+    sourceBySlug: new Map(),
+    trace: { passes: [] },
+    cost: { ms: 0 },
+    failureReason: null,
+  });
+});
+
+afterEach(() => {
+  mock.restore();
+});
+
+describe("memory-v3 shadow middleware", () => {
+  test("flag off → byte-for-byte pass-through, no v3 call, no log write", async () => {
+    v3Enabled = false;
+    v3Shadow = false;
+    let nextCalls = 0;
+    const args = makeArgs();
+    const result = await memoryV3ShadowMiddleware(
+      args,
+      async (a) => {
+        nextCalls++;
+        // identity is preserved — pass-through hands the same args down.
+        expect(a).toBe(args);
+        return DOWNSTREAM_RESULT;
+      },
+      makeCtx(),
+    );
+
+    // Returns the exact downstream object reference, unchanged.
+    expect(result).toBe(DOWNSTREAM_RESULT);
+    expect(nextCalls).toBe(1);
+
+    await flush();
+    expect(loopCalls.length).toBe(0);
+    expect(logCalls.length).toBe(0);
+  });
+
+  test("enabled but shadow off → still a pure pass-through", async () => {
+    v3Enabled = true;
+    v3Shadow = false;
+    const args = makeArgs();
+    const result = await memoryV3ShadowMiddleware(
+      args,
+      async () => DOWNSTREAM_RESULT,
+      makeCtx(),
+    );
+    expect(result).toBe(DOWNSTREAM_RESULT);
+    await flush();
+    expect(loopCalls.length).toBe(0);
+    expect(logCalls.length).toBe(0);
+  });
+
+  test("v3 config block absent → pass-through, no throw, no v3 call", async () => {
+    // Reproduces the agent-loop test fixtures (and any config built outside the
+    // Zod schema) where `memory.v3` is undefined. The gate must not throw.
+    v3Present = false;
+    const args = makeArgs();
+    const result = await memoryV3ShadowMiddleware(
+      args,
+      async () => DOWNSTREAM_RESULT,
+      makeCtx(),
+    );
+    expect(result).toBe(DOWNSTREAM_RESULT);
+    await flush();
+    expect(loopCalls.length).toBe(0);
+    expect(logCalls.length).toBe(0);
+  });
+
+  test("flag on → v3 runs, v3_shadow row logged, downstream result unchanged", async () => {
+    v3Enabled = true;
+    v3Shadow = true;
+    loopImpl = async () => ({
+      selectedSlugs: ["topic/a", "topic/b"],
+      sourceBySlug: new Map([["topic/a", "dense"]]),
+      trace: { passes: [] },
+      cost: { ms: 12 },
+      failureReason: null,
+    });
+
+    const args = makeArgs();
+    const result = await memoryV3ShadowMiddleware(
+      args,
+      async () => DOWNSTREAM_RESULT,
+      makeCtx(),
+    );
+
+    // The injected result is the v2/default result, NOT v3.
+    expect(result).toBe(DOWNSTREAM_RESULT);
+
+    await flush();
+
+    // v3 ran exactly once, with a faithfully-built RetrievalInput.
+    expect(loopCalls.length).toBe(1);
+    const input = loopCalls[0]!.input;
+    expect(input.nowText).toBe("NOW context");
+    expect(input.workspaceDir).toBe("/tmp/shadow-test-workspace");
+    expect(input.priorEverInjected).toEqual([{ slug: "old/page", turn: 0 }]);
+    expect(input.recentTurnPairs.length).toBeGreaterThan(0);
+    expect(input.recentTurnPairs.at(-1)?.userMessage).toBe("hello memory");
+
+    // Exactly one v3_shadow row, carrying v3's selection.
+    expect(logCalls.length).toBe(1);
+    const logged = logCalls[0]!;
+    expect(logged.mode).toBe("v3_shadow");
+    expect(logged.conversationId).toBe("conv-shadow");
+    expect(logged.turn).toBe(3);
+    expect(logged.concepts.map((c) => c.slug)).toEqual(["topic/a", "topic/b"]);
+
+    // Per-slug lane provenance is carried from sourceBySlug; a slug absent from
+    // the map (topic/b) gets no lane rather than a bogus one.
+    const bySlug = new Map(logged.concepts.map((c) => [c.slug, c]));
+    expect(bySlug.get("topic/a")?.lane).toBe("dense");
+    expect(bySlug.get("topic/b")?.lane).toBeUndefined();
+  });
+
+  test("v3 error → logged/swallowed, turn result unaffected, no log row", async () => {
+    v3Enabled = true;
+    v3Shadow = true;
+    loopImpl = async () => {
+      throw new Error("v3 boom");
+    };
+
+    const args = makeArgs();
+    // The middleware must not reject even though the detached shadow throws.
+    const result = await memoryV3ShadowMiddleware(
+      args,
+      async () => DOWNSTREAM_RESULT,
+      makeCtx(),
+    );
+    expect(result).toBe(DOWNSTREAM_RESULT);
+
+    await flush();
+    // Loop was attempted; the failure short-circuited before logging.
+    expect(loopCalls.length).toBe(1);
+    expect(logCalls.length).toBe(0);
+  });
+
+  test("aborted signal → shadow does no v3 work", async () => {
+    v3Enabled = true;
+    v3Shadow = true;
+    const controller = new AbortController();
+    controller.abort();
+
+    const result = await memoryV3ShadowMiddleware(
+      makeArgs(controller.signal),
+      async () => DOWNSTREAM_RESULT,
+      makeCtx(),
+    );
+    expect(result).toBe(DOWNSTREAM_RESULT);
+
+    await flush();
+    expect(loopCalls.length).toBe(0);
+    expect(logCalls.length).toBe(0);
+  });
+
+  test("untrusted actor → no v3 work, even with shadow on", async () => {
+    v3Enabled = true;
+    v3Shadow = true;
+
+    for (const trustClass of ["trusted_contact", "unknown"] as const) {
+      loopCalls.length = 0;
+      logCalls.length = 0;
+      const untrusted: TrustContext = { sourceChannel: "vellum", trustClass };
+      const result = await memoryV3ShadowMiddleware(
+        makeArgs(undefined, untrusted),
+        async () => DOWNSTREAM_RESULT,
+        makeCtx(),
+      );
+      expect(result).toBe(DOWNSTREAM_RESULT);
+
+      await flush();
+      // Untrusted turn never spends shadow retrieval LLM calls.
+      expect(loopCalls.length).toBe(0);
+      expect(logCalls.length).toBe(0);
+    }
+  });
+
+  test("absent trust context → treated as untrusted, no v3 work", async () => {
+    v3Enabled = true;
+    v3Shadow = true;
+
+    // Build args with an explicitly-undefined trustContext (the `makeArgs`
+    // default would substitute the guardian fixture, so construct directly).
+    const args: MemoryArgs = {
+      conversationId: "conv-shadow",
+      trustContext: undefined,
+      turnIndex: 3,
+      signal: new AbortController().signal,
+    };
+    const result = await memoryV3ShadowMiddleware(
+      args,
+      async () => DOWNSTREAM_RESULT,
+      makeCtx(),
+    );
+    expect(result).toBe(DOWNSTREAM_RESULT);
+
+    await flush();
+    expect(loopCalls.length).toBe(0);
+    expect(logCalls.length).toBe(0);
+  });
+
+  test("guardian actor → shadow runs (trust gate lets trusted turns through)", async () => {
+    v3Enabled = true;
+    v3Shadow = true;
+    const guardian: TrustContext = {
+      sourceChannel: "vellum",
+      trustClass: "guardian",
+    };
+
+    const result = await memoryV3ShadowMiddleware(
+      makeArgs(undefined, guardian),
+      async () => DOWNSTREAM_RESULT,
+      makeCtx(),
+    );
+    expect(result).toBe(DOWNSTREAM_RESULT);
+
+    await flush();
+    // Trusted turn fires the v3 loop and logs its selection.
+    expect(loopCalls.length).toBe(1);
+    expect(logCalls.length).toBe(1);
+  });
+});