Skip to content

[AI-assisted] Fix memory_grep default session recall #303

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
IWhatsskill wants to merge 2 commits into from
Closed

[AI-assisted] Fix memory_grep default session recall #303

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
71b8a19
fix: include default session collection in memory_grep
IWhatsskill Jun 2, 2026
ec85f53
fix: tighten memory_grep result classification
IWhatsskill Jun 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
153 changes: 117 additions & 36 deletions src/tools/memory-recall.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,13 @@ const MAX_GREP_RESULTS = 50;
const MAX_GREP_CHARS = 40000;
const MAX_SNIPPET_CHARS = 200;

type GrepSearchResult = {
id: string;
score: number;
text: string;
metadataJson?: Uint8Array;
};

// ── Schemas ──

const MEMORY_DESCRIBE_SCHEMA = {
Expand Down Expand Up @@ -183,6 +190,85 @@ function safeMatch(text: string, pattern: string, mode: "regex" | "text"): boole
}
}

function uniqueCollections(collections: string[]): string[] {
return [...new Set(collections.filter((collection) => collection.length > 0))];
}

function buildGrepCollections(sessionId: string, scope: "messages" | "summaries" | "both"): string[] {
const collections: string[] = [];
if (scope === "summaries" || scope === "both") {
collections.push(`session_summary:${sessionId}`);
}
if (scope === "messages" || scope === "both") {
collections.push(`session_raw:${sessionId}`);
}
if (sessionId.length > 0) {
// Keep memory_grep aligned with memory_search's default session recall collection.
// Older or default runtimes store session hits under session:<id> rather than
// the experimental session_summary/session_raw split.
collections.push(`session:${sessionId}`);
}
return uniqueCollections(collections);
}

async function searchGrepCollections(
client: Awaited<ReturnType<ClientGetter>>,
collections: string[],
pattern: string,
k: number,
): Promise<GrepSearchResult[]> {
if (collections.length === 0) {
return [];
}
if (collections.length === 1) {
const result = await client.searchText({
collection: collections[0],
text: pattern,
k,
});
return result.results ?? [];
}
const result = await client.searchTextCollections({
collections,
text: pattern,
k,
excludeByCollection: {},
});
return result.results ?? [];
}

function parseResultMetadata(result: GrepSearchResult): Record<string, unknown> {
if (result.metadataJson && result.metadataJson.length > 0) {
try {
return JSON.parse(new TextDecoder().decode(result.metadataJson)) as Record<string, unknown>;
} catch {
return {};
}
}
return {};
}

function resultCollection(meta: Record<string, unknown>, fallback = ""): string {
return typeof meta.collection === "string" ? meta.collection : fallback;
}

function isSummaryResult(result: GrepSearchResult, meta: Record<string, unknown>, collection: string): boolean {
return (
collection.startsWith("session_summary:") ||
result.id.startsWith("sum_") ||
typeof meta.eviction_cue === "string" ||
typeof meta.compaction_generation === "number"
);
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

function isTurnResult(_result: GrepSearchResult, meta: Record<string, unknown>, collection: string): boolean {
return (
collection.startsWith("session_raw:") ||
collection.startsWith("session:") ||
typeof meta.role === "string"
);
}

// ── Tool factories ──

export function createMemoryDescribeTool(
Expand Down Expand Up @@ -403,51 +489,46 @@ export function createMemoryGrepTool(
let totalChars = 0;
let totalMatches = 0;

if (scope === "summaries" || scope === "both") {
const searchK = Math.min(limit * 3, 200);
const summaryResults = await client.searchText({
collection: `session_summary:${sessionId}`,
text: pattern,
k: searchK,
});
for (const r of (summaryResults.results ?? [])) {
if (summaries.length >= limit || totalChars >= MAX_GREP_CHARS) break;
const searchK = Math.min(limit * 3, 200);
const grepResults = await searchGrepCollections(
client,
buildGrepCollections(sessionId, scope),
pattern,
searchK,
);
const seen = new Set<string>();

for (const r of grepResults) {
if (
(summaries.length >= limit && turns.length >= limit) ||
totalChars >= MAX_GREP_CHARS
) {
break;
}
const meta = parseResultMetadata(r);
const collection = resultCollection(meta);
const dedupKey = r.id;
if (seen.has(dedupKey)) continue;
seen.add(dedupKey);
const summaryResult = isSummaryResult(r, meta, collection);
const turnResult = isTurnResult(r, meta, collection);

if ((scope === "summaries" || scope === "both") && summaryResult) {
if (summaries.length >= limit) continue;
if (!safeMatch(r.text, pattern, mode)) continue;
totalMatches++;
let evictionCue: string | undefined;
if (r.metadataJson && r.metadataJson.length > 0) {
try {
const decoder = new TextDecoder();
const meta = JSON.parse(decoder.decode(r.metadataJson)) as Record<string, unknown>;
evictionCue = typeof meta.eviction_cue === "string" ? meta.eviction_cue : undefined;
} catch { /* best-effort */ }
}
const evictionCue = typeof meta.eviction_cue === "string" ? meta.eviction_cue : undefined;
const snippet = truncateSnippet(r.text);
summaries.push({ summaryId: r.id, snippet, score: r.score, evictionCue });
totalChars += snippet.length;
continue;
}
}

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 ?? [])) {
if (turns.length >= limit || totalChars >= MAX_GREP_CHARS) break;
if ((scope === "messages" || scope === "both") && (turnResult || !summaryResult)) {
if (turns.length >= limit) continue;
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 */ }
}
const role = typeof meta.role === "string" ? meta.role : "unknown";
turns.push({ turnId: r.id, snippet, role, score: r.score });
totalChars += snippet.length;
}
Expand Down
107 changes: 105 additions & 2 deletions test/unit/memory-recall.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,58 @@ class FakeRecallClient {
}],
};
}

async searchTextCollections(params: Record<string, unknown>) {
this.calls.push({ method: "searchTextCollections", params });
const collections = params.collections as string[] | undefined;
return {
results: [{
id: "sum-1",
score: 0.9,
text: "needle inside summary text",
metadataJson: new TextEncoder().encode(JSON.stringify({
collection: collections?.[0] ?? "session_summary:active-session",
role: "assistant",
eviction_cue: "summary cue",
})),
}],
};
}
}

class DefaultSessionRecallClient extends FakeRecallClient {
constructor(private readonly includeDuplicateRawHit = false) {
super();
}

override async searchTextCollections(params: Record<string, unknown>) {
this.calls.push({ method: "searchTextCollections", params });
const duplicateRawHit = this.includeDuplicateRawHit
? [{
id: "turn-1",
score: 0.71,
text: "needle inside duplicate raw collection",
metadataJson: new TextEncoder().encode(JSON.stringify({
collection: "session_raw:active-session",
role: "user",
})),
}]
: [];
return {
results: [
...duplicateRawHit,
{
id: "turn-1",
score: 0.88,
text: "needle inside default session collection",
metadataJson: new TextEncoder().encode(JSON.stringify({
collection: "session:active-session",
role: "user",
})),
},
],
};
}
}

function fakeRuntime(client: FakeRecallClient): PluginRuntime {
Expand Down Expand Up @@ -82,8 +134,59 @@ test("memory_grep defaults to the active session id", async () => {
const result = await tool.execute("call-1", { pattern: "needle", scope: "summaries" });

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[0]?.method, "searchTextCollections");
assert.deepEqual(client.calls[0]?.params.collections, [
"session_summary:active-session",
"session:active-session",
]);
});

test("memory_grep searches the default active session collection", async () => {
const client = new DefaultSessionRecallClient();
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; role: string }> };

assert.equal(details.totalMatches, 1);
assert.deepEqual(client.calls[0]?.params.collections, [
"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 deduplicates the same turn across raw and default session collections", async () => {
const client = new DefaultSessionRecallClient(true);
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; role: string }> };

assert.equal(details.totalMatches, 1);
assert.deepEqual(client.calls[0]?.params.collections, [
"session_raw:active-session",
"session:active-session",
]);
assert.deepEqual(details.turns, [{
turnId: "turn-1",
snippet: "needle inside duplicate raw collection",
role: "user",
score: 0.71,
}]);
});

test("memory_expand defaults to the active session id", async () => {
Expand Down