diff --git a/src/tools/memory-recall.ts b/src/tools/memory-recall.ts index 1b1439d4..959cdbb9 100644 --- a/src/tools/memory-recall.ts +++ b/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 { + 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; + } + } 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(); + 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; - 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; } } diff --git a/test/unit/memory-recall.test.ts b/test/unit/memory-recall.test.ts index 2bfffb99..64e5141d 100644 --- a/test/unit/memory-recall.test.ts +++ b/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 }> = []; @@ -29,7 +36,7 @@ class FakeRecallClient { }; } - async searchText(params: Record) { + async searchText(params: Record): Promise<{ results: FakeSearchResult[] }> { this.calls.push({ method: "searchText", params }); return { results: [{ @@ -45,6 +52,24 @@ class FakeRecallClient { } } +function encodeMetadata(value: Record): Uint8Array { + return new TextEncoder().encode(JSON.stringify(value)); +} + +class CollectionRecallClient extends FakeRecallClient { + constructor(private readonly resultsByCollection: Record) { + super(); + } + + override async searchText(params: Record): 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 () => {