[AI-assisted] Fix memory_grep default session recall

src/tools/memory-recall.ts

@@ -49,6 +49,13 @@ type MemoryGrepDetails = {
   truncated: boolean;
 };
 
+type GrepSearchResult = {
+  id: string;
+  score: number;
+  text: string;
+  metadataJson?: Uint8Array;
+};
+
 // ── Constants ──
 
 const MAX_EXPAND_TOKENS = 8000;
@@ -183,6 +190,30 @@ function safeMatch(text: string, pattern: string, mode: "regex" | "text"): boole
   }
 }
 
+function parseGrepMetadata(result: GrepSearchResult): Record<string, unknown> {
+  if (result.metadataJson && result.metadataJson.length > 0) {
+    try {
+      const parsed = JSON.parse(new TextDecoder().decode(result.metadataJson)) as unknown;
+      if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) {
+        return parsed as Record<string, unknown>;
+      }
+    } catch {
+      return {};
+    }
+  }
+  return {};
+}
+
+function readGrepRole(result: GrepSearchResult): string {
+  const meta = parseGrepMetadata(result);
+  return typeof meta.role === "string" ? meta.role : "unknown";
+}
+
+function buildMessageGrepCollections(sessionId: string): string[] {
+  if (sessionId.length === 0) return [];
+  return [`session_raw:${sessionId}`, `session:${sessionId}`];
+}
+
 // ── Tool factories ──
 
 export function createMemoryDescribeTool(
@@ -432,26 +463,38 @@ export function createMemoryGrepTool(
 
         if (scope === "messages" || scope === "both") {
           const searchK = Math.min(limit * 3, 200);
-          const turnResults = await client.searchText({
-            collection: `session_raw:${sessionId}`,
-            text: pattern,
-            k: searchK,
-          });
-          for (const r of (turnResults.results ?? [])) {
+          const bestTurnById = new Map<string, {
+            turnId: string;
+            snippet: string;
+            role: string;
+            score: number;
+          }>();
+          for (const collection of buildMessageGrepCollections(sessionId)) {
+            const turnResults = await client.searchText({
+              collection,
+              text: pattern,
+              k: searchK,
+            });
+            for (const r of (turnResults.results ?? [])) {
+              if (!safeMatch(r.text, pattern, mode)) continue;
+              const candidate = {
+                turnId: r.id,
+                snippet: truncateSnippet(r.text),
+                role: readGrepRole(r),
+                score: r.score,
+              };
+              const existing = bestTurnById.get(r.id);
+              if (!existing || candidate.score > existing.score) {
+                bestTurnById.set(r.id, candidate);
+              }
+            }
+          }
+          const dedupedTurns = [...bestTurnById.values()].sort((a, b) => b.score - a.score);
+          for (const turn of dedupedTurns) {
             if (turns.length >= limit || totalChars >= MAX_GREP_CHARS) break;
-            if (!safeMatch(r.text, pattern, mode)) continue;
             totalMatches++;
-            const snippet = truncateSnippet(r.text);
-            let role = "unknown";
-            if (r.metadataJson && r.metadataJson.length > 0) {
-              try {
-                const decoder = new TextDecoder();
-                const meta = JSON.parse(decoder.decode(r.metadataJson)) as Record<string, unknown>;
-                role = typeof meta.role === "string" ? meta.role : "unknown";
-              } catch { /* best-effort */ }
-            }
-            turns.push({ turnId: r.id, snippet, role, score: r.score });
-            totalChars += snippet.length;
+            turns.push(turn);
+            totalChars += turn.snippet.length;
           }
         }
 

test/unit/memory-recall.test.ts

@@ -12,6 +12,13 @@ const silentLogger = {
   info(_message: string) {},
 };
 
+type FakeSearchResult = {
+  id: string;
+  score: number;
+  text: string;
+  metadataJson?: Uint8Array;
+};
+
 class FakeRecallClient {
   public calls: Array<{ method: string; params: Record<string, unknown> }> = [];
 
@@ -29,7 +36,7 @@ class FakeRecallClient {
     };
   }
 
-  async searchText(params: Record<string, unknown>) {
+  async searchText(params: Record<string, unknown>): Promise<{ results: FakeSearchResult[] }> {
     this.calls.push({ method: "searchText", params });
     return {
       results: [{
@@ -45,6 +52,24 @@ class FakeRecallClient {
   }
 }
 
+function encodeMetadata(value: Record<string, unknown>): Uint8Array {
+  return new TextEncoder().encode(JSON.stringify(value));
+}
+
+class CollectionRecallClient extends FakeRecallClient {
+  constructor(private readonly resultsByCollection: Record<string, FakeSearchResult[]>) {
+    super();
+  }
+
+  override async searchText(params: Record<string, unknown>): Promise<{ results: FakeSearchResult[] }> {
+    this.calls.push({ method: "searchText", params });
+    const collection = typeof params.collection === "string" ? params.collection : "";
+    return {
+      results: this.resultsByCollection[collection] ?? [],
+    };
+  }
+}
+
 function fakeRuntime(client: FakeRecallClient): PluginRuntime {
   return {
     getClient: async () => client as unknown as LibravDBClient,
@@ -84,6 +109,205 @@ test("memory_grep defaults to the active session id", async () => {
   assert.equal((result.details as { totalMatches: number }).totalMatches, 1);
   assert.equal(client.calls[0]?.method, "searchText");
   assert.equal(client.calls[0]?.params.collection, "session_summary:active-session");
+  assert.equal(client.calls.length, 1);
+});
+
+test("memory_grep does not query message collections without a session id", async () => {
+  const client = new CollectionRecallClient({
+    "session_raw:": [{
+      id: "turn-1",
+      score: 0.99,
+      text: "needle in malformed collection",
+      metadataJson: encodeMetadata({ role: "user" }),
+    }],
+  });
+  const tool = createMemoryGrepTool(
+    async () => client as unknown as LibravDBClient,
+    () => undefined,
+    silentLogger,
+  );
+
+  const result = await tool.execute("call-1", { pattern: "needle", scope: "messages" });
+  const details = result.details as { totalMatches: number; turns: Array<{ turnId: string }> };
+
+  assert.equal(details.totalMatches, 0);
+  assert.deepEqual(details.turns, []);
+  assert.equal(client.calls.length, 0);
+});
+
+test("memory_grep searches the default active session collection for messages", async () => {
+  const client = new CollectionRecallClient({
+    "session_raw:active-session": [],
+    "session:active-session": [{
+      id: "turn-1",
+      score: 0.88,
+      text: "needle inside default session collection",
+      metadataJson: encodeMetadata({ role: "user" }),
+    }],
+  });
+  const tool = createMemoryGrepTool(
+    async () => client as unknown as LibravDBClient,
+    () => "active-session",
+    silentLogger,
+  );
+
+  const result = await tool.execute("call-1", { pattern: "needle", scope: "messages" });
+  const details = result.details as { totalMatches: number; turns: Array<{ turnId: string; snippet: string; role: string; score: number }> };
+
+  assert.equal(details.totalMatches, 1);
+  assert.deepEqual(client.calls.map((call) => call.params.collection), [
+    "session_raw:active-session",
+    "session:active-session",
+  ]);
+  assert.deepEqual(details.turns, [{
+    turnId: "turn-1",
+    snippet: "needle inside default session collection",
+    role: "user",
+    score: 0.88,
+  }]);
+});
+
+test("memory_grep treats non-object message metadata as missing", async () => {
+  const client = new CollectionRecallClient({
+    "session_raw:active-session": [{
+      id: "turn-1",
+      score: 0.77,
+      text: "needle survives non-object metadata",
+      metadataJson: new TextEncoder().encode("null"),
+    }],
+    "session:active-session": [],
+  });
+  const tool = createMemoryGrepTool(
+    async () => client as unknown as LibravDBClient,
+    () => "active-session",
+    silentLogger,
+  );
+
+  const result = await tool.execute("call-1", { pattern: "needle", scope: "messages" });
+  const details = result.details as { totalMatches: number; turns: Array<{ turnId: string; snippet: string; role: string; score: number }> };
+
+  assert.equal(details.totalMatches, 1);
+  assert.deepEqual(details.turns, [{
+    turnId: "turn-1",
+    snippet: "needle survives non-object metadata",
+    role: "unknown",
+    score: 0.77,
+  }]);
+});
+
+test("memory_grep keeps independent summary and message budgets", async () => {
+  const client = new CollectionRecallClient({
+    "session_summary:active-session": [{
+      id: "sum_1",
+      score: 0.7,
+      text: "needle inside summary text",
+      metadataJson: encodeMetadata({ eviction_cue: "summary cue" }),
+    }],
+    "session_raw:active-session": [],
+    "session:active-session": [{
+      id: "turn-1",
+      score: 0.99,
+      text: "needle inside default session collection",
+      metadataJson: encodeMetadata({ role: "assistant" }),
+    }],
+  });
+  const tool = createMemoryGrepTool(
+    async () => client as unknown as LibravDBClient,
+    () => "active-session",
+    silentLogger,
+  );
+
+  const result = await tool.execute("call-1", { pattern: "needle", scope: "both", limit: 1 });
+  const details = result.details as {
+    totalMatches: number;
+    summaries: Array<{ summaryId: string; snippet: string; score: number; evictionCue?: string }>;
+    turns: Array<{ turnId: string; snippet: string; role: string; score: number }>;
+  };
+
+  assert.deepEqual(client.calls.map((call) => call.params.collection), [
+    "session_summary:active-session",
+    "session_raw:active-session",
+    "session:active-session",
+  ]);
+  assert.equal(details.totalMatches, 2);
+  assert.deepEqual(details.summaries, [{
+    summaryId: "sum_1",
+    snippet: "needle inside summary text",
+    score: 0.7,
+    evictionCue: "summary cue",
+  }]);
+  assert.deepEqual(details.turns, [{
+    turnId: "turn-1",
+    snippet: "needle inside default session collection",
+    role: "assistant",
+    score: 0.99,
+  }]);
+});
+
+test("memory_grep deduplicates messages only after exact matching", async () => {
+  const client = new CollectionRecallClient({
+    "session_raw:active-session": [{
+      id: "turn-1",
+      score: 0.99,
+      text: "semantic neighbor without the target phrase",
+      metadataJson: encodeMetadata({ role: "user" }),
+    }],
+    "session:active-session": [{
+      id: "turn-1",
+      score: 0.5,
+      text: "needle appears in the default session collection",
+      metadataJson: encodeMetadata({ role: "assistant" }),
+    }],
+  });
+  const tool = createMemoryGrepTool(
+    async () => client as unknown as LibravDBClient,
+    () => "active-session",
+    silentLogger,
+  );
+
+  const result = await tool.execute("call-1", { pattern: "needle", scope: "messages" });
+  const details = result.details as { totalMatches: number; turns: Array<{ turnId: string; snippet: string; role: string; score: number }> };
+
+  assert.equal(details.totalMatches, 1);
+  assert.deepEqual(details.turns, [{
+    turnId: "turn-1",
+    snippet: "needle appears in the default session collection",
+    role: "assistant",
+    score: 0.5,
+  }]);
+});
+
+test("memory_grep keeps the highest-scored duplicate message hit", async () => {
+  const client = new CollectionRecallClient({
+    "session_raw:active-session": [{
+      id: "turn-1",
+      score: 0.71,
+      text: "needle inside duplicate raw collection",
+      metadataJson: encodeMetadata({ role: "user" }),
+    }],
+    "session:active-session": [{
+      id: "turn-1",
+      score: 0.88,
+      text: "needle inside default session collection",
+      metadataJson: encodeMetadata({ role: "user" }),
+    }],
+  });
+  const tool = createMemoryGrepTool(
+    async () => client as unknown as LibravDBClient,
+    () => "active-session",
+    silentLogger,
+  );
+
+  const result = await tool.execute("call-1", { pattern: "needle", scope: "messages" });
+  const details = result.details as { totalMatches: number; turns: Array<{ turnId: string; snippet: string; role: string; score: number }> };
+
+  assert.equal(details.totalMatches, 1);
+  assert.deepEqual(details.turns, [{
+    turnId: "turn-1",
+    snippet: "needle inside default session collection",
+    role: "user",
+    score: 0.88,
+  }]);
 });
 
 test("memory_expand defaults to the active session id", async () => {