From a5b4f6a93210e748c40e64160848dac5c5fa94f1 Mon Sep 17 00:00:00 2001 From: clawsweeper <274271284+clawsweeper[bot]@users.noreply.github.com> Date: Mon, 4 May 2026 14:34:59 +0000 Subject: [PATCH] fix(memory): prevent all-corpus memory hit starvation --- .../memory-core/src/tools.citations.test.ts | 16 ++--- extensions/memory-core/src/tools.ts | 72 ++++++++++++++----- extensions/memory-wiki/src/query.test.ts | 56 +++++++++++++++ extensions/memory-wiki/src/query.ts | 51 ++++++++++--- 4 files changed, 158 insertions(+), 37 deletions(-) diff --git a/extensions/memory-core/src/tools.citations.test.ts b/extensions/memory-core/src/tools.citations.test.ts index 8c357070e92..3b58d95ad15 100644 --- a/extensions/memory-core/src/tools.citations.test.ts +++ b/extensions/memory-core/src/tools.citations.test.ts @@ -294,14 +294,6 @@ describe("memory tools", () => { snippet: "Memory result A", source: "memory" as const, }, - { - path: "memory/note-b.md", - startLine: 1, - endLine: 2, - score: 0.8, - snippet: "Memory result B", - source: "memory" as const, - }, ]); registerMemoryCorpusSupplement("memory-wiki", { search: async () => [ @@ -358,10 +350,14 @@ describe("memory tools", () => { const details = result.details as { results: Array<{ corpus: string; path: string }> }; const corpora = details.results.map((r) => r.corpus); - // Memory results must appear despite lower numeric scores. + // Memory results must appear despite lower numeric scores, and the spare + // memory quota should be backfilled by the remaining wiki result. expect(corpora).toContain("memory"); expect(corpora).toContain("wiki"); - expect(details.results.length).toBeLessThanOrEqual(5); + expect(details.results).toHaveLength(5); + expect( + details.results.filter((entry) => entry.corpus === "wiki").map((entry) => entry.path), + ).toEqual(["w1.md", "w2.md", "w3.md", "w4.md"]); }); it("merges memory and wiki corpus search results for corpus=all", async () => { diff --git a/extensions/memory-core/src/tools.ts b/extensions/memory-core/src/tools.ts index df78272b0a8..56216bf932e 100644 --- a/extensions/memory-core/src/tools.ts +++ b/extensions/memory-core/src/tools.ts @@ -5,6 +5,7 @@ import { jsonResult, readNumberParam, readStringParam, + type MemoryCorpusSearchResult, type OpenClawConfig, } from "openclaw/plugin-sdk/memory-core-host-runtime-core"; import type { @@ -35,6 +36,50 @@ import { searchMemoryCorpusSupplements, } from "./tools.shared.js"; +type MemorySearchToolResult = + | (Record & { corpus: "memory"; score: number; path: string }) + | MemoryCorpusSearchResult; + +function sortMemorySearchToolResults(results: T[]): T[] { + return results.toSorted((left, right) => { + if (left.score !== right.score) { + return right.score - left.score; + } + return left.path.localeCompare(right.path); + }); +} + +function mergeMemorySearchCorpusResults(params: { + memoryResults: MemorySearchToolResult[]; + supplementResults: MemorySearchToolResult[]; + maxResults: number; + balanceCorpora: boolean; +}): MemorySearchToolResult[] { + const memoryResults = sortMemorySearchToolResults(params.memoryResults); + const supplementResults = sortMemorySearchToolResults(params.supplementResults); + if (!params.balanceCorpora || memoryResults.length === 0 || supplementResults.length === 0) { + return sortMemorySearchToolResults([...memoryResults, ...supplementResults]).slice( + 0, + params.maxResults, + ); + } + + const perCorpusCap = Math.ceil(params.maxResults / 2); + const selectedMemory = memoryResults.slice(0, perCorpusCap); + const selectedSupplements = supplementResults.slice(0, perCorpusCap); + const selected = [...selectedMemory, ...selectedSupplements]; + if (selected.length < params.maxResults) { + selected.push( + ...sortMemorySearchToolResults([ + ...memoryResults.slice(selectedMemory.length), + ...supplementResults.slice(selectedSupplements.length), + ]).slice(0, params.maxResults - selected.length), + ); + } + + return sortMemorySearchToolResults(selected).slice(0, params.maxResults); +} + function buildRecallKey( result: Pick, ): string { @@ -319,26 +364,15 @@ export function createMemorySearchTool(options: { corpus: requestedCorpus, }) : []; - // When both corpora are present (corpus=all), wiki and memory scores - // use incomparable scales (integer point vs cosine similarity), so a - // raw-score sort lets wiki hits starve memory hits (#77337). - // Cap each corpus at ceil(maxResults/2) before the joint sort so both - // corpora are guaranteed representation in the final slice. + // Wiki and memory scores use incomparable scales, so corpus=all first + // balances candidate selection and then backfills any unused slots. const effectiveMax = Math.max(1, maxResults ?? 10); - const perCorpusCap = - requestedCorpus === "all" && supplementResults.length > 0 - ? Math.ceil(effectiveMax / 2) - : effectiveMax; - const cappedMemory = surfacedMemoryResults.slice(0, perCorpusCap); - const cappedSupplements = supplementResults.slice(0, perCorpusCap); - const results = [...cappedMemory, ...cappedSupplements] - .toSorted((left, right) => { - if (left.score !== right.score) { - return right.score - left.score; - } - return left.path.localeCompare(right.path); - }) - .slice(0, effectiveMax); + const results = mergeMemorySearchCorpusResults({ + memoryResults: surfacedMemoryResults, + supplementResults, + maxResults: effectiveMax, + balanceCorpora: requestedCorpus === "all", + }); return jsonResult({ results, provider, diff --git a/extensions/memory-wiki/src/query.test.ts b/extensions/memory-wiki/src/query.test.ts index 15131431db6..51ee87eeebf 100644 --- a/extensions/memory-wiki/src/query.test.ts +++ b/extensions/memory-wiki/src/query.test.ts @@ -578,6 +578,62 @@ describe("searchMemoryWiki", () => { }); }); + it("includes memory results and backfills wiki capacity for all-corpus search", async () => { + const { rootDir, config } = await createQueryVault({ + initialize: true, + config: { + search: { backend: "shared", corpus: "all" }, + }, + }); + for (const index of [1, 2, 3, 4, 5]) { + await fs.writeFile( + path.join(rootDir, "entities", `alpha-${index}.md`), + renderWikiMarkdown({ + frontmatter: { + pageType: "entity", + id: `entity.alpha.${index}`, + title: `Alpha ${index}`, + }, + body: `# Alpha ${index}\n\nalpha wiki ${index}\n`, + }), + "utf8", + ); + } + const manager = createMemoryManager({ + searchResults: [ + { + path: "MEMORY.md", + startLine: 4, + endLine: 8, + score: 0.9, + snippet: "alpha durable memory", + source: "memory", + citation: "MEMORY.md#L4-L8", + }, + ], + }); + getActiveMemorySearchManagerMock.mockResolvedValue({ manager }); + + const results = await searchMemoryWiki({ + config, + appConfig: createAppConfig(), + query: "alpha", + maxResults: 5, + }); + + expect(results).toHaveLength(5); + expect(results.some((result) => result.corpus === "memory")).toBe(true); + expect( + results.filter((result) => result.corpus === "wiki").map((result) => result.path), + ).toEqual([ + "entities/alpha-1.md", + "entities/alpha-2.md", + "entities/alpha-3.md", + "entities/alpha-4.md", + ]); + expect(manager.search).toHaveBeenCalledWith("alpha", { maxResults: 5 }); + }); + it("uses the active session agent for shared memory search", async () => { const { config } = await createQueryVault({ initialize: true, diff --git a/extensions/memory-wiki/src/query.ts b/extensions/memory-wiki/src/query.ts index 720e6e7a38b..0dbd5680714 100644 --- a/extensions/memory-wiki/src/query.ts +++ b/extensions/memory-wiki/src/query.ts @@ -183,6 +183,43 @@ type QuerySearchOverrides = { searchCorpus?: WikiSearchCorpus; }; +function sortWikiSearchResults(results: WikiSearchResult[]): WikiSearchResult[] { + return results.toSorted((left, right) => { + if (left.score !== right.score) { + return right.score - left.score; + } + return left.title.localeCompare(right.title); + }); +} + +function mergeWikiSearchCorpusResults(params: { + wikiResults: WikiSearchResult[]; + memoryResults: WikiSearchResult[]; + maxResults: number; + balanceCorpora: boolean; +}): WikiSearchResult[] { + const wikiResults = sortWikiSearchResults(params.wikiResults); + const memoryResults = sortWikiSearchResults(params.memoryResults); + if (!params.balanceCorpora || wikiResults.length === 0 || memoryResults.length === 0) { + return sortWikiSearchResults([...wikiResults, ...memoryResults]).slice(0, params.maxResults); + } + + const perCorpusCap = Math.ceil(params.maxResults / 2); + const selectedWiki = wikiResults.slice(0, perCorpusCap); + const selectedMemory = memoryResults.slice(0, perCorpusCap); + const selected = [...selectedWiki, ...selectedMemory]; + if (selected.length < params.maxResults) { + selected.push( + ...sortWikiSearchResults([ + ...wikiResults.slice(selectedWiki.length), + ...memoryResults.slice(selectedMemory.length), + ]).slice(0, params.maxResults - selected.length), + ); + } + + return sortWikiSearchResults(selected).slice(0, params.maxResults); +} + async function listWikiMarkdownFiles(rootDir: string): Promise { const files = ( await Promise.all( @@ -1219,14 +1256,12 @@ export async function searchMemoryWiki(params: { ) : []; - return [...wikiResults, ...memoryResults] - .toSorted((left, right) => { - if (left.score !== right.score) { - return right.score - left.score; - } - return left.title.localeCompare(right.title); - }) - .slice(0, maxResults); + return mergeWikiSearchCorpusResults({ + wikiResults, + memoryResults, + maxResults, + balanceCorpora: effectiveConfig.search.corpus === "all", + }); } export async function getMemoryWikiPage(params: {