@vellumai/assistant 0.8.6 → 0.8.7-dev.202606052118.34cd356

package/src/memory/graph/__tests__/conversation-graph-memory-registry.test.ts

@@ -0,0 +1,119 @@
+/**
+ * Tests for the live, per-conversation {@link ConversationGraphMemory} registry
+ * (`getLiveGraphMemory`). The registry lets memory-domain code that only knows
+ * a conversation id — notably the post-compaction re-injection hook — reach the
+ * same in-memory handle the turn's retrieval mutated, without the agent loop
+ * threading the handle through its generic context.
+ */
+import { describe, expect, mock, test } from "bun:test";
+
+import { createMockLoggerModule } from "../../../__tests__/helpers/mock-logger.js";
+
+mock.module("../../../util/logger.js", () => createMockLoggerModule());
+
+const { ConversationGraphMemory, getLiveGraphMemory } =
+  await import("../conversation-graph-memory.js");
+
+describe("ConversationGraphMemory live registry", () => {
+  test("a constructed handle is discoverable by its conversation id", () => {
+    /**
+     * Tests that constructing a handle registers it so it can be looked up.
+     */
+
+    // GIVEN a conversation id
+    const conversationId = `conv-registry-${crypto.randomUUID()}`;
+
+    // WHEN a graph handle is constructed for it
+    const handle = new ConversationGraphMemory(conversationId);
+
+    // THEN the registry returns that exact instance
+    expect(getLiveGraphMemory(conversationId)).toBe(handle);
+  });
+
+  test("dispose removes the handle from the registry", () => {
+    /**
+     * Tests that disposing a handle unregisters it.
+     */
+
+    // GIVEN a registered graph handle
+    const conversationId = `conv-registry-${crypto.randomUUID()}`;
+    const handle = new ConversationGraphMemory(conversationId);
+
+    // WHEN it is disposed
+    handle.dispose();
+
+    // THEN the registry no longer resolves the conversation id
+    expect(getLiveGraphMemory(conversationId)).toBeUndefined();
+  });
+
+  test("missing id and undefined id both resolve to undefined", () => {
+    /**
+     * Tests the lookup's absence handling for unknown and undefined keys.
+     */
+
+    // GIVEN no handle registered for these keys
+    // WHEN the registry is queried with an unknown id and with undefined
+    // THEN both resolve to undefined (the hook treats absence as "no graph")
+    expect(
+      getLiveGraphMemory(`conv-never-${crypto.randomUUID()}`),
+    ).toBeUndefined();
+    expect(getLiveGraphMemory(undefined)).toBeUndefined();
+  });
+
+  test("recreating for the same id replaces the registered handle", () => {
+    /**
+     * Tests that the latest constructed handle wins for a conversation id.
+     */
+
+    // GIVEN a handle already registered for a conversation id
+    const conversationId = `conv-registry-${crypto.randomUUID()}`;
+    new ConversationGraphMemory(conversationId);
+
+    // WHEN a second handle is constructed for the same id
+    const recreated = new ConversationGraphMemory(conversationId);
+
+    // THEN the registry resolves to the most recent instance
+    expect(getLiveGraphMemory(conversationId)).toBe(recreated);
+  });
+
+  test("disposing a stale handle does not evict the current one", () => {
+    /**
+     * Tests the dispose guard during an eviction + recreation race: a stale
+     * handle must not delete the live entry that now points at a newer handle.
+     */
+
+    // GIVEN a superseded (stale) handle and the current handle for one id
+    const conversationId = `conv-registry-${crypto.randomUUID()}`;
+    const stale = new ConversationGraphMemory(conversationId);
+    const current = new ConversationGraphMemory(conversationId);
+
+    // WHEN the stale handle is disposed
+    stale.dispose();
+
+    // THEN the registry still resolves to the current handle
+    expect(getLiveGraphMemory(conversationId)).toBe(current);
+  });
+
+  test("records and exposes the PKB query vector pair from the registry", () => {
+    /**
+     * Tests that the dense/sparse PKB query pair recorded during retrieval is
+     * readable off the same live handle the PKB-reminder injector looks up.
+     */
+
+    // GIVEN a registered graph handle with no recorded vectors yet
+    const conversationId = `conv-registry-${crypto.randomUUID()}`;
+    const handle = new ConversationGraphMemory(conversationId);
+    expect(handle.pkbQueryVector).toBeUndefined();
+    expect(handle.pkbSparseVector).toBeUndefined();
+
+    // WHEN a retrieval records the turn's dense/sparse pair
+    const dense = [0.1, 0.2, 0.3];
+    const sparse = { indices: [0, 2], values: [0.5, 0.9] };
+    handle.recordPkbQueryVectors(dense, sparse);
+
+    // THEN a registry consumer reads back the same pair by conversation id
+    const looked = getLiveGraphMemory(conversationId);
+    expect(looked?.pkbQueryVector).toBe(dense);
+    expect(looked?.pkbSparseVector).toBe(sparse);
+  });
+});

package/src/memory/graph/consolidation.ts

@@ -279,17 +279,14 @@ async function identifyDuplicateGroups(
 
   const systemPrompt = `You are scanning a list of memory nodes for DUPLICATES — nodes that describe the exact same specific event or fact. Group duplicates together. Two nodes are duplicates ONLY if they describe the same underlying thing with substantially the same details. Be conservative — nodes about the same person or topic but with different details, timestamps, or context are NOT duplicates. Only include nodes that have at least one true duplicate.`;
 
-  const response = await provider.sendMessage(
-    [userMessage(listing)],
-    [DUPE_DETECT_TOOL],
+  const response = await provider.sendMessage([userMessage(listing)], {
+    tools: [DUPE_DETECT_TOOL],
     systemPrompt,
-    {
-      config: {
-        callSite: "memoryConsolidation" as const,
-        tool_choice: { type: "tool" as const, name: "report_duplicate_groups" },
-      },
+    config: {
+      callSite: "memoryConsolidation" as const,
+      tool_choice: { type: "tool" as const, name: "report_duplicate_groups" },
     },
-  );
+  });
 
   const toolBlock = extractToolUse(response);
   if (!toolBlock) return [];
@@ -458,9 +455,9 @@ async function consolidateChunk(
         "Consolidate this partition. Focus on merging duplicates and fading old memories.",
       ),
     ],
-    [CONSOLIDATE_TOOL_SCHEMA],
-    systemPrompt,
     {
+      tools: [CONSOLIDATE_TOOL_SCHEMA],
+      systemPrompt,
       config: {
         callSite: "memoryConsolidation" as const,
         tool_choice: { type: "tool" as const, name: "consolidate_diff" },

package/src/memory/graph/conversation-graph-memory.ts

@@ -22,6 +22,7 @@ import { getWorkspaceDir } from "../../util/platform.js";
 import { getDb } from "../db-connection.js";
 import { embedWithRetry } from "../embed.js";
 import { generateSparseEmbedding } from "../embedding-backend.js";
+import { wrapMemoryBlock } from "../memory-marker.js";
 import type { QdrantSparseVector } from "../qdrant-client.js";
 import { memorySummaries } from "../schema.js";
 import { conversations } from "../schema/conversations.js";
@@ -61,6 +62,31 @@ const ESTIMATED_IMAGE_TOKENS = 1000;
 // Per-conversation state
 // ---------------------------------------------------------------------------
 
+/**
+ * Registry of the live, per-conversation graph handles keyed by conversation
+ * id. A handle registers itself on construction and removes itself on
+ * {@link ConversationGraphMemory.dispose}, so memory-domain code that only
+ * knows a conversation id (e.g. the post-compaction re-injection hook) can
+ * reach the same in-memory handle the turn's retrieval used — its live
+ * `tracker` / cached-node state, which a DB-reconstructed handle would not
+ * carry. Not a general service locator: it holds only the graph handle, and
+ * the daemon's `Conversation` remains the owner of the instance's lifecycle.
+ */
+const liveByConversation = new Map<string, ConversationGraphMemory>();
+
+/**
+ * Look up the live {@link ConversationGraphMemory} for a conversation, or
+ * `undefined` when none is registered (no active conversation, or a context
+ * with no conversation id). Returns the same instance the turn's retrieval
+ * mutated, so cached-node re-tracking operates on real state.
+ */
+export function getLiveGraphMemory(
+  conversationId: string | undefined,
+): ConversationGraphMemory | undefined {
+  if (!conversationId) return undefined;
+  return liveByConversation.get(conversationId);
+}
+
 /**
  * Manages memory graph state for a single conversation.
  * Create one per Conversation instance. Persists across turns.
@@ -74,9 +100,24 @@ export class ConversationGraphMemory {
   private lastInjectedBlock: string | null = null;
   private lastInjectedNodeIds: string[] = [];
   private lastInjectedImages: Map<string, ResolvedImage> = new Map();
+  private lastPkbQueryVector: number[] | undefined;
+  private lastPkbSparseVector: QdrantSparseVector | undefined;
 
   constructor(conversationId: string) {
     this.conversationId = conversationId;
+    liveByConversation.set(conversationId, this);
+  }
+
+  /**
+   * Remove this handle from the live registry. Called from
+   * `Conversation.dispose`. Guards against clobbering a newer handle for the
+   * same conversation (eviction + recreation) by only deleting the entry when
+   * it still points at this instance.
+   */
+  dispose(): void {
+    if (liveByConversation.get(this.conversationId) === this) {
+      liveByConversation.delete(this.conversationId);
+    }
   }
 
   /**
@@ -304,6 +345,31 @@ export class ConversationGraphMemory {
     this.tracker.add(this.lastInjectedNodeIds);
   }
 
+  /**
+   * Record the dense/sparse query-vector pair this turn's retrieval produced
+   * for PKB hybrid search. The PKB-reminder injector reuses the same
+   * embedding (looked up by conversation id via {@link getLiveGraphMemory})
+   * rather than receiving it threaded through the agent loop, so the vectors
+   * stay owned by the memory-retrieval domain that computes them.
+   */
+  recordPkbQueryVectors(
+    dense: number[] | undefined,
+    sparse: QdrantSparseVector | undefined,
+  ): void {
+    this.lastPkbQueryVector = dense;
+    this.lastPkbSparseVector = sparse;
+  }
+
+  /** Dense PKB query vector from this turn's retrieval, or `undefined`. */
+  get pkbQueryVector(): number[] | undefined {
+    return this.lastPkbQueryVector;
+  }
+
+  /** Sparse PKB query vector paired with {@link pkbQueryVector}. */
+  get pkbSparseVector(): QdrantSparseVector | undefined {
+    return this.lastPkbSparseVector;
+  }
+
   /**
    * Main entry point — called on every turn before the LLM sees the messages.
    *
@@ -895,13 +961,45 @@ export function stripExistingMemoryInjections(messages: Message[]): Message[] {
   const last = messages[messages.length - 1];
   if (!last || last.role !== "user") return messages;
 
-  const firstNonMemory = countMemoryPrefixBlocks(last.content);
-  if (firstNonMemory === 0) return messages;
+  const stripped = stripMemoryPrefixFromUserMessage(last);
+  if (stripped === last) return messages;
 
-  return [
-    ...messages.slice(0, -1),
-    { ...last, content: last.content.slice(firstNonMemory) },
-  ];
+  return [...messages.slice(0, -1), stripped];
+}
+
+/**
+ * Strip the memory-injected prefix from a single user message. Returns the
+ * same message reference unchanged when it is not a user message or carries no
+ * injected prefix, so callers can cheaply detect no-ops. Shared by
+ * `stripExistingMemoryInjections` (last message only) and
+ * `stripAllMemoryInjections` (every user message).
+ */
+function stripMemoryPrefixFromUserMessage(message: Message): Message {
+  if (message.role !== "user") return message;
+  const firstNonMemory = countMemoryPrefixBlocks(message.content);
+  if (firstNonMemory === 0) return message;
+  return { ...message, content: message.content.slice(firstNonMemory) };
+}
+
+/**
+ * Remove memory-injected blocks from EVERY user message, not just the last.
+ *
+ * memory-v3 live mode strips the injected `<memory>` block from all historical
+ * user messages each turn: it keeps the prompt lean and makes history
+ * byte-stable for prompt caching (v2 strips only the last user message — see
+ * `stripExistingMemoryInjections`). Reuses the same per-message block
+ * recognition as the last-message strip. Wired into the runtime-assembly
+ * suppression step (`conversation-runtime-assembly.ts`), which calls it when
+ * `memory-v3-live` is on and the v3 injector produced a block this turn.
+ */
+export function stripAllMemoryInjections(messages: Message[]): Message[] {
+  let changed = false;
+  const result = messages.map((message) => {
+    const stripped = stripMemoryPrefixFromUserMessage(message);
+    if (stripped !== message) changed = true;
+    return stripped;
+  });
+  return changed ? result : messages;
 }
 
 /**
@@ -933,7 +1031,7 @@ function injectTextBlock(messages: Message[], text: string): Message[] {
       content: [
         {
           type: "text" as const,
-          text: `<memory>\n${text}\n</memory>`,
+          text: wrapMemoryBlock(text),
         },
         ...userTail.content,
       ],
@@ -977,7 +1075,7 @@ function injectMemoryBlock(
 
   blocks.push({
     type: "text" as const,
-    text: `<memory>\n${text}\n</memory>`,
+    text: wrapMemoryBlock(text),
   });
 
   return [

package/src/memory/graph/extraction.ts

@@ -1071,17 +1071,14 @@ export async function runGraphExtraction(
         ),
       ];
 
-  const response = await provider.sendMessage(
-    extractionMessages,
-    [EXTRACT_TOOL_SCHEMA],
+  const response = await provider.sendMessage(extractionMessages, {
+    tools: [EXTRACT_TOOL_SCHEMA],
     systemPrompt,
-    {
-      config: {
-        callSite: "memoryExtraction" as const,
-        tool_choice: { type: "tool" as const, name: "extract_graph_diff" },
-      },
+    config: {
+      callSite: "memoryExtraction" as const,
+      tool_choice: { type: "tool" as const, name: "extract_graph_diff" },
     },
-  );
+  });
 
   const toolBlock = extractToolUse(response);
   if (!toolBlock) {

package/src/memory/graph/narrative.ts

@@ -192,9 +192,9 @@ export async function runNarrativeRefinement(
         "Review and refine the narrative structure of these memories. Identify story arcs and assign roles with the benefit of hindsight.",
       ),
     ],
-    [NARRATIVE_TOOL_SCHEMA],
-    systemPrompt,
     {
+      tools: [NARRATIVE_TOOL_SCHEMA],
+      systemPrompt,
       config: {
         callSite: "narrativeRefinement" as const,
         tool_choice: { type: "tool" as const, name: "refine_narratives" },

package/src/memory/graph/pattern-scan.ts

@@ -161,9 +161,9 @@ export async function runPatternScan(
         "Analyze this memory sample for recurring patterns. Only report patterns you're confident about.",
       ),
     ],
-    [PATTERN_TOOL_SCHEMA],
-    systemPrompt,
     {
+      tools: [PATTERN_TOOL_SCHEMA],
+      systemPrompt,
       config: {
         callSite: "patternScan" as const,
         tool_choice: { type: "tool" as const, name: "detect_patterns" },

package/src/memory/graph/retriever.ts

@@ -113,10 +113,9 @@ async function rerankAndDedup(
       })
       .join("\n");
 
-    const response = await provider.sendMessage(
-      [userMessage(listing)],
-      [RERANK_TOOL],
-      `You are selecting memories for an AI assistant's context at conversation start. You see ${candidates.length} candidate memories ranked by algorithmic score.
+    const response = await provider.sendMessage([userMessage(listing)], {
+      tools: [RERANK_TOOL],
+      systemPrompt: `You are selecting memories for an AI assistant's context at conversation start. You see ${candidates.length} candidate memories ranked by algorithmic score.
 
 Your job:
 1. REMOVE DUPLICATES: If multiple entries describe the same event/fact/topic, keep ONLY the most complete version. Be aggressive — even partial overlaps should be deduplicated.
@@ -125,15 +124,13 @@ Your job:
    - Diversity (don't load 5 memories about the same topic)
    - Importance (key relationship moments, active commitments, identity-defining events)
 3. Return the IDs in order of importance (most important first).`,
-      {
-        config: {
-          callSite: "memoryRetrieval" as const,
-          tool_choice: { type: "tool" as const, name: "select_memories" },
-          thinking: { type: "disabled" },
-          temperature: 0,
-        },
+      config: {
+        callSite: "memoryRetrieval" as const,
+        tool_choice: { type: "tool" as const, name: "select_memories" },
+        thinking: { type: "disabled" },
+        temperature: 0,
       },
-    );
+    });
 
     const toolBlock = extractToolUse(response);
     if (!toolBlock) return candidates.slice(0, maxNodes);
@@ -214,9 +211,9 @@ async function dedupForTurn(
 
     const response = await provider.sendMessage(
       [userMessage(`query:\n${query}\n\nitems:\n\n${listing}`)],
-      [SELECT_ITEMS_TOOL],
-      `Dedupe + rerank the following numbered items. Pick the most relevant items to the query. Call the select_items tool.\n\nBe aggressive on dedup — when multiple items describe the same event, fact, or status, keep ONLY the richest version. But be generous on relevance — only cut items that are completely irrelevant to the query. If it's even tangentially related, keep it.`,
       {
+        tools: [SELECT_ITEMS_TOOL],
+        systemPrompt: `Dedupe + rerank the following numbered items. Pick the most relevant items to the query. Call the select_items tool.\n\nBe aggressive on dedup — when multiple items describe the same event, fact, or status, keep ONLY the richest version. But be generous on relevance — only cut items that are completely irrelevant to the query. If it's even tangentially related, keep it.`,
         config: {
           callSite: "memoryRetrieval" as const,
           tool_choice: { type: "tool" as const, name: "select_items" },
@@ -304,19 +301,16 @@ async function dedupCrossCategory(
       })
       .join("\n");
 
-    const response = await provider.sendMessage(
-      [userMessage(listing)],
-      [DEDUP_ITEMS_TOOL],
-      `Deduplicate the following numbered items. When multiple items describe the same event, fact, or status, keep ONLY the richest version. Keep ALL items that are not duplicates — do not filter by relevance or topic. Call the select_items tool with every item that survives dedup.`,
-      {
-        config: {
-          callSite: "memoryRetrieval" as const,
-          tool_choice: { type: "tool" as const, name: "select_items" },
-          thinking: { type: "disabled" },
-          temperature: 0,
-        },
+    const response = await provider.sendMessage([userMessage(listing)], {
+      tools: [DEDUP_ITEMS_TOOL],
+      systemPrompt: `Deduplicate the following numbered items. When multiple items describe the same event, fact, or status, keep ONLY the richest version. Keep ALL items that are not duplicates — do not filter by relevance or topic. Call the select_items tool with every item that survives dedup.`,
+      config: {
+        callSite: "memoryRetrieval" as const,
+        tool_choice: { type: "tool" as const, name: "select_items" },
+        thinking: { type: "disabled" },
+        temperature: 0,
       },
-    );
+    });
 
     const toolBlock = extractToolUse(response);
     if (!toolBlock) return candidates.slice(0, maxNodes);

package/src/memory/graph/tools.ts

@@ -15,7 +15,7 @@ const RECALL_DEPTHS = ["fast", "standard", "deep"] as const;
  * Explicit local information search across memory, conversations, and
  * workspace files.
  */
-export const graphRecallDefinition: ToolDefinition = {
+export const graphRecallDefinition = {
   name: "recall",
   description:
     'Search local information the moment you feel uncertain. Use recall for memory, past conversations, and workspace files — before you guess, before you ask, before you hedge. Auto-injection is incomplete by design; it surfaces patterns, not the specifics you need to answer well. If you catch yourself reaching for "I think", "I believe", "if I remember", "didn\'t we", "last time" — that\'s the signal. Recall. If a turn references someone, a place, a decision, a document, or prior work you should be able to find locally — recall. Call it multiple times per conversation if the turn warrants it. Be specific in your query for best results.',
@@ -51,7 +51,7 @@ export const graphRecallDefinition: ToolDefinition = {
     },
     required: ["query"],
   },
-};
+} satisfies ToolDefinition;
 
 /**
  * `remember` tool description. The retrospective pass catches what isn't
@@ -70,7 +70,7 @@ const REMEMBER_DESCRIPTION =
  * of the buffer into longer-form storage runs as a separate periodic job in
  * both modes.
  */
-export const graphRememberDefinition: ToolDefinition = {
+export const graphRememberDefinition = {
   name: "remember",
   description: REMEMBER_DESCRIPTION,
   input_schema: {
@@ -89,4 +89,4 @@ export const graphRememberDefinition: ToolDefinition = {
     },
     required: ["content"],
   },
-};
+} satisfies ToolDefinition;

package/src/memory/job-handlers/conversation-starters.ts

@@ -22,6 +22,7 @@ import {
   checkpointKey,
   CK_BATCH,
   CK_ITEM_COUNT,
+  CK_LAST_ATTEMPT_AT,
   CK_LAST_GEN_AT,
   countActiveMemoryNodes,
   getCheckpointValue,
@@ -295,47 +296,48 @@ Bad → Good (incomplete phrase → complete):
           "Generate personalized conversation starters based on my context.",
         ),
       ],
-      [
-        {
-          name: "store_conversation_starters",
-          description: "Store generated conversation starter suggestions",
-          input_schema: {
-            type: "object" as const,
-            properties: {
-              starters: {
-                type: "array",
-                items: {
-                  type: "object",
-                  properties: {
-                    label: {
-                      type: "string",
-                      description:
-                        "User-voice chip label (2-7 words, max 40 chars, verb-first)",
-                    },
-                    prompt: {
-                      type: "string",
-                      description:
-                        "Full message sent on click (1-2 natural sentences, as the user would say it)",
-                    },
-                    category: {
-                      type: "string",
-                      enum: [...CONVERSATION_STARTER_CATEGORIES],
-                      description: "Capability category for grouping",
+      {
+        tools: [
+          {
+            name: "store_conversation_starters",
+            description: "Store generated conversation starter suggestions",
+            input_schema: {
+              type: "object" as const,
+              properties: {
+                starters: {
+                  type: "array",
+                  items: {
+                    type: "object",
+                    properties: {
+                      label: {
+                        type: "string",
+                        description:
+                          "User-voice chip label (2-7 words, max 40 chars, verb-first)",
+                      },
+                      prompt: {
+                        type: "string",
+                        description:
+                          "Full message sent on click (1-2 natural sentences, as the user would say it)",
+                      },
+                      category: {
+                        type: "string",
+                        enum: [...CONVERSATION_STARTER_CATEGORIES],
+                        description: "Capability category for grouping",
+                      },
                     },
+                    required: ["label", "prompt", "category"],
                   },
-                  required: ["label", "prompt", "category"],
                 },
               },
+              required: ["starters"],
             },
-            required: ["starters"],
           },
-        },
-      ],
-      systemPrompt,
-      {
+        ],
+        systemPrompt,
         config: {
           callSite: "conversationStarters" as const,
           max_tokens: 2048,
+          temperature: 0.7,
           tool_choice: {
             type: "tool" as const,
             name: "store_conversation_starters",
@@ -395,18 +397,27 @@ export async function generateConversationStartersJob(
   const db = getDb();
   const now = Date.now();
 
+  // Record attempt time so the route handler's cooldown prevents rapid retries
+  // even when generation produces 0 valid starters.
+  upsertCheckpoint(
+    checkpointKey(CK_LAST_ATTEMPT_AT, scopeId),
+    String(now),
+    now,
+  );
+
   const starters = await generateStarters(scopeId);
   if (starters.length === 0) {
     log.info({ scopeId }, "No conversation starters generated");
 
-    // Sync checkpoints so both `staleByAge` and `checkpointAhead` clear.
+    // Sync the item count checkpoint so `checkpointAhead` clears, but do NOT
+    // update `CK_LAST_GEN_AT` — the stale TTL should trigger a retry on the
+    // next GET rather than blocking retries for the full TTL window.
     const totalActive = countActiveMemoryNodes(scopeId);
     upsertCheckpoint(
       checkpointKey(CK_ITEM_COUNT, scopeId),
       String(totalActive),
       now,
     );
-    upsertCheckpoint(checkpointKey(CK_LAST_GEN_AT, scopeId), String(now), now);
     return;
   }
 

package/src/memory/job-handlers/summarization.ts

@@ -180,9 +180,8 @@ async function summarizeWithLLM(
     try {
       const response = await provider.sendMessage(
         [userMessage(userParts.join("\n"))],
-        undefined,
-        systemPrompt,
         {
+          systemPrompt,
           config: {
             callSite: "conversationSummarization" as const,
             max_tokens: SUMMARY_MAX_TOKENS,