Skip to content

Commit a5b4f6a

Browse files
clawsweeper[bot]clawsweeper[bot]
committed
fix(memory): prevent all-corpus memory hit starvation
1 parent 5add41c commit a5b4f6a

4 files changed

Lines changed: 158 additions & 37 deletions

File tree

‎extensions/memory-core/src/tools.citations.test.ts‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -294,14 +294,6 @@ describe("memory tools", () => {
294294
snippet: "Memory result A",
295295
source: "memory" as const,
296296
},
297-
{
298-
path: "memory/note-b.md",
299-
startLine: 1,
300-
endLine: 2,
301-
score: 0.8,
302-
snippet: "Memory result B",
303-
source: "memory" as const,
304-
},
305297
]);
306298
registerMemoryCorpusSupplement("memory-wiki", {
307299
search: async () => [
@@ -358,10 +350,14 @@ describe("memory tools", () => {
358350
const details = result.details as { results: Array<{ corpus: string; path: string }> };
359351
const corpora = details.results.map((r) => r.corpus);
360352

361-
// Memory results must appear despite lower numeric scores.
353+
// Memory results must appear despite lower numeric scores, and the spare
354+
// memory quota should be backfilled by the remaining wiki result.
362355
expect(corpora).toContain("memory");
363356
expect(corpora).toContain("wiki");
364-
expect(details.results.length).toBeLessThanOrEqual(5);
357+
expect(details.results).toHaveLength(5);
358+
expect(
359+
details.results.filter((entry) => entry.corpus === "wiki").map((entry) => entry.path),
360+
).toEqual(["w1.md", "w2.md", "w3.md", "w4.md"]);
365361
});
366362

367363
it("merges memory and wiki corpus search results for corpus=all", async () => {

‎extensions/memory-core/src/tools.ts‎

Lines changed: 53 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import {
55
jsonResult,
66
readNumberParam,
77
readStringParam,
8+
type MemoryCorpusSearchResult,
89
type OpenClawConfig,
910
} from "openclaw/plugin-sdk/memory-core-host-runtime-core";
1011
import type {
@@ -35,6 +36,50 @@ import {
3536
searchMemoryCorpusSupplements,
3637
} from "./tools.shared.js";
3738

39+
type MemorySearchToolResult =
40+
| (Record<string, unknown> & { corpus: "memory"; score: number; path: string })
41+
| MemoryCorpusSearchResult;
42+
43+
function sortMemorySearchToolResults<T extends { score: number; path: string }>(results: T[]): T[] {
44+
return results.toSorted((left, right) => {
45+
if (left.score !== right.score) {
46+
return right.score - left.score;
47+
}
48+
return left.path.localeCompare(right.path);
49+
});
50+
}
51+
52+
function mergeMemorySearchCorpusResults(params: {
53+
memoryResults: MemorySearchToolResult[];
54+
supplementResults: MemorySearchToolResult[];
55+
maxResults: number;
56+
balanceCorpora: boolean;
57+
}): MemorySearchToolResult[] {
58+
const memoryResults = sortMemorySearchToolResults(params.memoryResults);
59+
const supplementResults = sortMemorySearchToolResults(params.supplementResults);
60+
if (!params.balanceCorpora || memoryResults.length === 0 || supplementResults.length === 0) {
61+
return sortMemorySearchToolResults([...memoryResults, ...supplementResults]).slice(
62+
0,
63+
params.maxResults,
64+
);
65+
}
66+
67+
const perCorpusCap = Math.ceil(params.maxResults / 2);
68+
const selectedMemory = memoryResults.slice(0, perCorpusCap);
69+
const selectedSupplements = supplementResults.slice(0, perCorpusCap);
70+
const selected = [...selectedMemory, ...selectedSupplements];
71+
if (selected.length < params.maxResults) {
72+
selected.push(
73+
...sortMemorySearchToolResults([
74+
...memoryResults.slice(selectedMemory.length),
75+
...supplementResults.slice(selectedSupplements.length),
76+
]).slice(0, params.maxResults - selected.length),
77+
);
78+
}
79+
80+
return sortMemorySearchToolResults(selected).slice(0, params.maxResults);
81+
}
82+
3883
function buildRecallKey(
3984
result: Pick<MemorySearchResult, "source" | "path" | "startLine" | "endLine">,
4085
): string {
@@ -319,26 +364,15 @@ export function createMemorySearchTool(options: {
319364
corpus: requestedCorpus,
320365
})
321366
: [];
322-
// When both corpora are present (corpus=all), wiki and memory scores
323-
// use incomparable scales (integer point vs cosine similarity), so a
324-
// raw-score sort lets wiki hits starve memory hits (#77337).
325-
// Cap each corpus at ceil(maxResults/2) before the joint sort so both
326-
// corpora are guaranteed representation in the final slice.
367+
// Wiki and memory scores use incomparable scales, so corpus=all first
368+
// balances candidate selection and then backfills any unused slots.
327369
const effectiveMax = Math.max(1, maxResults ?? 10);
328-
const perCorpusCap =
329-
requestedCorpus === "all" && supplementResults.length > 0
330-
? Math.ceil(effectiveMax / 2)
331-
: effectiveMax;
332-
const cappedMemory = surfacedMemoryResults.slice(0, perCorpusCap);
333-
const cappedSupplements = supplementResults.slice(0, perCorpusCap);
334-
const results = [...cappedMemory, ...cappedSupplements]
335-
.toSorted((left, right) => {
336-
if (left.score !== right.score) {
337-
return right.score - left.score;
338-
}
339-
return left.path.localeCompare(right.path);
340-
})
341-
.slice(0, effectiveMax);
370+
const results = mergeMemorySearchCorpusResults({
371+
memoryResults: surfacedMemoryResults,
372+
supplementResults,
373+
maxResults: effectiveMax,
374+
balanceCorpora: requestedCorpus === "all",
375+
});
342376
return jsonResult({
343377
results,
344378
provider,

‎extensions/memory-wiki/src/query.test.ts‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -578,6 +578,62 @@ describe("searchMemoryWiki", () => {
578578
});
579579
});
580580

581+
it("includes memory results and backfills wiki capacity for all-corpus search", async () => {
582+
const { rootDir, config } = await createQueryVault({
583+
initialize: true,
584+
config: {
585+
search: { backend: "shared", corpus: "all" },
586+
},
587+
});
588+
for (const index of [1, 2, 3, 4, 5]) {
589+
await fs.writeFile(
590+
path.join(rootDir, "entities", `alpha-${index}.md`),
591+
renderWikiMarkdown({
592+
frontmatter: {
593+
pageType: "entity",
594+
id: `entity.alpha.${index}`,
595+
title: `Alpha ${index}`,
596+
},
597+
body: `# Alpha ${index}\n\nalpha wiki ${index}\n`,
598+
}),
599+
"utf8",
600+
);
601+
}
602+
const manager = createMemoryManager({
603+
searchResults: [
604+
{
605+
path: "MEMORY.md",
606+
startLine: 4,
607+
endLine: 8,
608+
score: 0.9,
609+
snippet: "alpha durable memory",
610+
source: "memory",
611+
citation: "MEMORY.md#L4-L8",
612+
},
613+
],
614+
});
615+
getActiveMemorySearchManagerMock.mockResolvedValue({ manager });
616+
617+
const results = await searchMemoryWiki({
618+
config,
619+
appConfig: createAppConfig(),
620+
query: "alpha",
621+
maxResults: 5,
622+
});
623+
624+
expect(results).toHaveLength(5);
625+
expect(results.some((result) => result.corpus === "memory")).toBe(true);
626+
expect(
627+
results.filter((result) => result.corpus === "wiki").map((result) => result.path),
628+
).toEqual([
629+
"entities/alpha-1.md",
630+
"entities/alpha-2.md",
631+
"entities/alpha-3.md",
632+
"entities/alpha-4.md",
633+
]);
634+
expect(manager.search).toHaveBeenCalledWith("alpha", { maxResults: 5 });
635+
});
636+
581637
it("uses the active session agent for shared memory search", async () => {
582638
const { config } = await createQueryVault({
583639
initialize: true,

‎extensions/memory-wiki/src/query.ts‎

Lines changed: 43 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,43 @@ type QuerySearchOverrides = {
183183
searchCorpus?: WikiSearchCorpus;
184184
};
185185

186+
function sortWikiSearchResults(results: WikiSearchResult[]): WikiSearchResult[] {
187+
return results.toSorted((left, right) => {
188+
if (left.score !== right.score) {
189+
return right.score - left.score;
190+
}
191+
return left.title.localeCompare(right.title);
192+
});
193+
}
194+
195+
function mergeWikiSearchCorpusResults(params: {
196+
wikiResults: WikiSearchResult[];
197+
memoryResults: WikiSearchResult[];
198+
maxResults: number;
199+
balanceCorpora: boolean;
200+
}): WikiSearchResult[] {
201+
const wikiResults = sortWikiSearchResults(params.wikiResults);
202+
const memoryResults = sortWikiSearchResults(params.memoryResults);
203+
if (!params.balanceCorpora || wikiResults.length === 0 || memoryResults.length === 0) {
204+
return sortWikiSearchResults([...wikiResults, ...memoryResults]).slice(0, params.maxResults);
205+
}
206+
207+
const perCorpusCap = Math.ceil(params.maxResults / 2);
208+
const selectedWiki = wikiResults.slice(0, perCorpusCap);
209+
const selectedMemory = memoryResults.slice(0, perCorpusCap);
210+
const selected = [...selectedWiki, ...selectedMemory];
211+
if (selected.length < params.maxResults) {
212+
selected.push(
213+
...sortWikiSearchResults([
214+
...wikiResults.slice(selectedWiki.length),
215+
...memoryResults.slice(selectedMemory.length),
216+
]).slice(0, params.maxResults - selected.length),
217+
);
218+
}
219+
220+
return sortWikiSearchResults(selected).slice(0, params.maxResults);
221+
}
222+
186223
async function listWikiMarkdownFiles(rootDir: string): Promise<string[]> {
187224
const files = (
188225
await Promise.all(
@@ -1219,14 +1256,12 @@ export async function searchMemoryWiki(params: {
12191256
)
12201257
: [];
12211258

1222-
return [...wikiResults, ...memoryResults]
1223-
.toSorted((left, right) => {
1224-
if (left.score !== right.score) {
1225-
return right.score - left.score;
1226-
}
1227-
return left.title.localeCompare(right.title);
1228-
})
1229-
.slice(0, maxResults);
1259+
return mergeWikiSearchCorpusResults({
1260+
wikiResults,
1261+
memoryResults,
1262+
maxResults,
1263+
balanceCorpora: effectiveConfig.search.corpus === "all",
1264+
});
12301265
}
12311266

12321267
export async function getMemoryWikiPage(params: {

0 commit comments

Comments
 (0)