From 3fe833dfa4084ad79d4545093f39f192c2108317 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Wed, 19 Aug 2026 20:57:32 -0700 Subject: [PATCH] fix(memory): preserve requested corpus outcomes (#126530) --- config/assertion-safety-baseline.txt | 2 +- extensions/memory-core/src/memory-corpus.ts | 225 +++++++ .../memory-core/src/memory-get-corpus.test.ts | 351 ++++++++++ .../memory-core/src/memory-read-tool.ts | 95 +++ .../src/memory-search-tool-query.ts | 44 +- .../memory-core/src/tools.citations.test.ts | 257 ++------ extensions/memory-core/src/tools.shared.ts | 58 +- extensions/memory-core/src/tools.ts | 609 +++++++----------- .../src/agent-vault-isolation.test.ts | 5 +- .../src/corpus-supplement.visibility.test.ts | 10 +- 10 files changed, 993 insertions(+), 663 deletions(-) create mode 100644 extensions/memory-core/src/memory-corpus.ts create mode 100644 extensions/memory-core/src/memory-get-corpus.test.ts create mode 100644 extensions/memory-core/src/memory-read-tool.ts diff --git a/config/assertion-safety-baseline.txt b/config/assertion-safety-baseline.txt index ca32eada005c..4f629fe6a098 100644 --- a/config/assertion-safety-baseline.txt +++ b/config/assertion-safety-baseline.txt @@ -803,7 +803,7 @@ extensions/memory-core/src/short-term-promotion-utils.ts 4 extensions/memory-core/src/short-term-promotion.ts 1 extensions/memory-core/src/standing-intents-tool.ts 5 extensions/memory-core/src/standing-intents.ts 2 -extensions/memory-core/src/tools.ts 4 +extensions/memory-core/src/tools.ts 2 extensions/memory-lancedb/config.ts 4 extensions/memory-lancedb/doctor-contract-api.ts 1 extensions/memory-lancedb/embeddings.ts 2 diff --git a/extensions/memory-core/src/memory-corpus.ts b/extensions/memory-core/src/memory-corpus.ts new file mode 100644 index 000000000000..f68c4bf2ec66 --- /dev/null +++ b/extensions/memory-core/src/memory-corpus.ts @@ -0,0 +1,225 @@ +import { runTasksWithConcurrency } from "openclaw/plugin-sdk/concurrency-runtime"; +import { formatErrorMessage } from "openclaw/plugin-sdk/error-runtime"; +import { + listMemoryCorpusSupplements, + type MemoryCorpusSearchResult, +} from "openclaw/plugin-sdk/memory-core-host-runtime-core"; +import { + DEFAULT_MEMORY_SEARCH_TIMEOUT_MS, + resolveMemorySearchAbortError, +} from "./memory/search-deadline.js"; + +type MemoryCorpus = "memory" | "wiki"; +type MemorySupplement = ReturnType[number]; +type MemorySupplementGetResult = NonNullable< + Awaited> +>; +type MemorySupplementReadResult = Omit & { + status: "ok"; + text: string; +}; + +export type MemoryCorpusAttempt = + | { corpus: MemoryCorpus; outcome: "ok"; value: T } + | { corpus: MemoryCorpus; outcome: "unavailable"; value: T; error: string } + | { corpus: MemoryCorpus; outcome: "not-registered" }; + +export function unavailableMemoryCorpus( + corpus: MemoryCorpus, + value: T, + error: unknown, +): MemoryCorpusAttempt { + return { corpus, outcome: "unavailable", value, error: formatErrorMessage(error) }; +} + +async function raceMemoryCorpusSignal(signal: AbortSignal, task: Promise): Promise { + if (signal.aborted) { + throw resolveMemorySearchAbortError(signal); + } + let removeAbort = () => {}; + const aborted = new Promise((_resolve, reject) => { + const onAbort = () => reject(resolveMemorySearchAbortError(signal)); + signal.addEventListener("abort", onAbort, { once: true }); + removeAbort = () => signal.removeEventListener("abort", onAbort); + }); + try { + return await Promise.race([task, aborted]); + } finally { + removeAbort(); + } +} + +export async function attemptMemoryCorpus(params: { + corpus: MemoryCorpus; + signal: AbortSignal; + unavailableValue: T; + run: () => Promise; +}): Promise> { + try { + return { + corpus: params.corpus, + outcome: "ok", + value: await raceMemoryCorpusSignal(params.signal, params.run()), + }; + } catch (error) { + return unavailableMemoryCorpus(params.corpus, params.unavailableValue, error); + } +} + +export async function runMemoryCorpusDeadline(params: { + operation: "memory_search" | "memory_get"; + parentSignal?: AbortSignal; + run: (signal: AbortSignal) => Promise; +}): Promise { + if (params.parentSignal?.aborted) { + throw resolveMemorySearchAbortError(params.parentSignal); + } + const controller = new AbortController(); + const timer = setTimeout(() => { + controller.abort( + new Error(`${params.operation} timed out after ${DEFAULT_MEMORY_SEARCH_TIMEOUT_MS / 1000}s`), + ); + }, DEFAULT_MEMORY_SEARCH_TIMEOUT_MS); + timer.unref?.(); + const onParentAbort = () => controller.abort(resolveMemorySearchAbortError(params.parentSignal!)); + params.parentSignal?.addEventListener("abort", onParentAbort, { once: true }); + try { + const result = await params.run(controller.signal); + if (params.parentSignal?.aborted) { + throw resolveMemorySearchAbortError(params.parentSignal); + } + return result; + } finally { + clearTimeout(timer); + params.parentSignal?.removeEventListener("abort", onParentAbort); + } +} + +export function composeMemoryCorpusMetadata( + attempts: readonly MemoryCorpusAttempt[], + extraWarnings: readonly string[] = [], +) { + const ordered = attempts.toSorted( + (left, right) => Number(left.corpus === "wiki") - Number(right.corpus === "wiki"), + ); + const warnings = ordered.flatMap((attempt) => { + if (attempt.outcome === "ok") { + return []; + } + const label = attempt.corpus === "memory" ? "Memory" : "Wiki"; + return [ + attempt.outcome === "not-registered" + ? `${label} corpus is not registered; results do not cover that requested corpus.` + : `${label} corpus unavailable: ${attempt.error}`, + ]; + }); + warnings.push(...extraWarnings); + const errors = ordered.flatMap((attempt) => + attempt.outcome === "unavailable" ? [attempt.error] : [], + ); + return { + corpora: ordered.map((attempt) => + attempt.outcome === "unavailable" + ? { corpus: attempt.corpus, outcome: attempt.outcome, error: attempt.error } + : { corpus: attempt.corpus, outcome: attempt.outcome }, + ), + ...(warnings.length > 0 ? { warning: warnings.join(" ") } : {}), + ...(errors.length > 0 ? { error: errors.join("; ") } : {}), + }; +} + +async function settleMemorySupplements(params: { + signal: AbortSignal; + run: (registration: MemorySupplement) => Promise; + merge: (results: T[]) => T; +}): Promise> { + const supplements = listMemoryCorpusSupplements().toSorted((left, right) => + left.pluginId.localeCompare(right.pluginId), + ); + if (supplements.length === 0) { + return { corpus: "wiki", outcome: "not-registered" }; + } + const failures: Array<{ pluginId: string; error: string }> = []; + const completed: Array = Array.from({ length: supplements.length }); + try { + await raceMemoryCorpusSignal( + params.signal, + runTasksWithConcurrency({ + tasks: supplements.map((registration, index) => async () => { + const result = await params.run(registration); + completed[index] = result; + return result; + }), + limit: Math.min(4, supplements.length), + onTaskError: (error, index) => { + failures.push({ + pluginId: supplements[index]!.pluginId, + error: formatErrorMessage(error), + }); + }, + }), + ); + } catch (error) { + return unavailableMemoryCorpus( + "wiki", + params.merge(completed.filter((result): result is T => result !== undefined)), + error, + ); + } + const value = params.merge(completed.filter((result): result is T => result !== undefined)); + if (failures.length === 0) { + return { corpus: "wiki", outcome: "ok", value }; + } + const orderedFailures = failures.toSorted((left, right) => + left.pluginId.localeCompare(right.pluginId), + ); + const error = + orderedFailures.length === 1 + ? orderedFailures[0]!.error + : orderedFailures.map((entry) => `${entry.pluginId}: ${entry.error}`).join("; "); + return { corpus: "wiki", outcome: "unavailable", value, error }; +} + +export async function searchMemoryCorpusSupplements(params: { + query: string; + maxResults?: number; + agentId?: string; + agentSessionKey?: string; + sandboxed?: boolean; + signal: AbortSignal; +}): Promise> { + const { signal, ...query } = params; + return await settleMemorySupplements({ + signal, + run: async ({ supplement }) => await supplement.search(query), + merge: (results) => + results + .flat() + .toSorted((left, right) => right.score - left.score || left.path.localeCompare(right.path)) + .slice(0, Math.max(1, params.maxResults ?? 10)), + }); +} + +export async function readMemoryCorpusSupplements(params: { + lookup: string; + fromLine?: number; + lineCount?: number; + agentId?: string; + agentSessionKey?: string; + sandboxed?: boolean; + signal: AbortSignal; +}): Promise> { + const { signal, ...query } = params; + return await settleMemorySupplements({ + signal, + run: async ({ supplement }) => { + const result = await supplement.get(query); + if (!result) { + return null; + } + const { content, ...details } = result; + return { ...details, status: "ok", text: content }; + }, + merge: (results) => results.find((result) => result !== null) ?? null, + }); +} diff --git a/extensions/memory-core/src/memory-get-corpus.test.ts b/extensions/memory-core/src/memory-get-corpus.test.ts new file mode 100644 index 000000000000..4282f65c7c1e --- /dev/null +++ b/extensions/memory-core/src/memory-get-corpus.test.ts @@ -0,0 +1,351 @@ +import type { MemoryReadResult } from "openclaw/plugin-sdk/memory-core-host-engine-storage"; +import { + clearMemoryPluginState, + registerMemoryCorpusSupplement, +} from "openclaw/plugin-sdk/memory-host-core"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { + resetMemoryToolMockState, + setMemoryReadFileImpl, +} from "./memory-tool-manager.test-mocks.js"; +import { createMemoryGetTool } from "./tools.js"; +import { asOpenClawConfig, createMemoryGetToolOrThrow } from "./tools.test-helpers.js"; + +const lookup = "memory/entities/alpha.md"; +const memoryHit = { + status: "ok", + text: "Memory entry", + path: lookup, + from: 1, + lines: 1, +} as const; +const emptyRange = { + status: "ok", + text: "", + path: lookup, + from: 10, + lines: 0, +} as const; +const memoryMiss = { status: "not_found", text: "", path: lookup } as const; +const wikiHit = { + corpus: "wiki", + path: lookup, + title: "Alpha", + kind: "entity", + content: "Wiki entry", + fromLine: 3, + lineCount: 5, +} as const; +const { content: wikiContent, ...wikiDetails } = wikiHit; +const wikiPayload = { ...wikiDetails, status: "ok" as const, text: wikiContent }; + +beforeEach(() => { + vi.useRealTimers(); + clearMemoryPluginState(); + resetMemoryToolMockState(); +}); + +type MemoryCase = MemoryReadResult | Error; +type WikiCase = "unregistered" | null | Error | typeof wikiHit; + +function prepareCase(memory: MemoryCase, wiki: WikiCase) { + setMemoryReadFileImpl(async () => { + if (memory instanceof Error) { + throw memory; + } + return memory; + }); + const get = vi.fn(async () => { + if (wiki instanceof Error) { + throw wiki; + } + return wiki === "unregistered" ? null : wiki; + }); + if (wiki !== "unregistered") { + registerMemoryCorpusSupplement("memory-wiki", { search: async () => [], get }); + } + return get; +} + +describe("memory_get corpus outcomes", () => { + it.each([ + { + name: "keeps a memory hit while reporting a successful wiki lookup", + memory: memoryHit, + wiki: wikiHit, + expected: { + ...memoryHit, + corpora: [ + { corpus: "memory", outcome: "ok" }, + { corpus: "wiki", outcome: "ok" }, + ], + }, + }, + { + name: "keeps a successful empty range while reporting wiki failure", + memory: emptyRange, + wiki: new Error("wiki unavailable"), + expected: { + ...emptyRange, + corpora: [ + { corpus: "memory", outcome: "ok" }, + { corpus: "wiki", outcome: "unavailable", error: "wiki unavailable" }, + ], + warning: "Wiki corpus unavailable: wiki unavailable", + error: "wiki unavailable", + }, + }, + { + name: "falls through from a memory miss to wiki content", + memory: memoryMiss, + wiki: wikiHit, + expected: { + ...wikiPayload, + corpora: [ + { corpus: "memory", outcome: "ok" }, + { corpus: "wiki", outcome: "ok" }, + ], + }, + }, + { + name: "keeps wiki content when memory is unavailable", + memory: new Error("memory unavailable"), + wiki: wikiHit, + expected: { + ...wikiPayload, + corpora: [ + { corpus: "memory", outcome: "unavailable", error: "memory unavailable" }, + { corpus: "wiki", outcome: "ok" }, + ], + warning: "Memory corpus unavailable: memory unavailable", + error: "memory unavailable", + }, + }, + { + name: "returns not found when every available corpus misses", + memory: memoryMiss, + wiki: null, + expected: { + ...memoryMiss, + corpora: [ + { corpus: "memory", outcome: "ok" }, + { corpus: "wiki", outcome: "ok" }, + ], + }, + }, + { + name: "returns not found when wiki misses and memory is unavailable", + memory: new Error("memory unavailable"), + wiki: null, + expected: { + ...memoryMiss, + corpora: [ + { corpus: "memory", outcome: "unavailable", error: "memory unavailable" }, + { corpus: "wiki", outcome: "ok" }, + ], + warning: "Memory corpus unavailable: memory unavailable", + error: "memory unavailable", + }, + }, + { + name: "keeps a proven miss when wiki is unavailable", + memory: memoryMiss, + wiki: new Error("wiki unavailable"), + expected: { + ...memoryMiss, + corpora: [ + { corpus: "memory", outcome: "ok" }, + { corpus: "wiki", outcome: "unavailable", error: "wiki unavailable" }, + ], + warning: "Wiki corpus unavailable: wiki unavailable", + error: "wiki unavailable", + }, + }, + { + name: "does not claim not found when no corpus was available", + memory: new Error("memory unavailable"), + wiki: new Error("wiki unavailable"), + expected: { + path: lookup, + text: "", + disabled: true, + corpora: [ + { corpus: "memory", outcome: "unavailable", error: "memory unavailable" }, + { corpus: "wiki", outcome: "unavailable", error: "wiki unavailable" }, + ], + warning: + "Memory corpus unavailable: memory unavailable Wiki corpus unavailable: wiki unavailable", + error: "memory unavailable; wiki unavailable", + }, + }, + { + name: "does not claim not found when memory failed and wiki is not registered", + memory: new Error("memory unavailable"), + wiki: "unregistered" as const, + expected: { + path: lookup, + text: "", + disabled: true, + corpora: [ + { corpus: "memory", outcome: "unavailable", error: "memory unavailable" }, + { corpus: "wiki", outcome: "not-registered" }, + ], + warning: + "Memory corpus unavailable: memory unavailable Wiki corpus is not registered; results do not cover that requested corpus.", + error: "memory unavailable", + }, + }, + ])("$name", async ({ memory, wiki, expected }) => { + const get = prepareCase(memory, wiki); + + const result = await createMemoryGetToolOrThrow().execute("call_get_all", { + path: lookup, + corpus: "all", + }); + + expect(result.details).toEqual(expected); + if (wiki !== "unregistered") { + expect(get).toHaveBeenCalledOnce(); + } + }); + + it.each([ + { + name: "registered miss", + wiki: null, + expected: { + ...memoryMiss, + corpora: [{ corpus: "wiki", outcome: "ok" }], + }, + }, + { + name: "unregistered corpus", + wiki: "unregistered" as const, + expected: { + path: lookup, + text: "", + corpora: [{ corpus: "wiki", outcome: "not-registered" }], + warning: "Wiki corpus is not registered; results do not cover that requested corpus.", + }, + }, + { + name: "unavailable corpus", + wiki: new Error("wiki unavailable"), + expected: { + path: lookup, + text: "", + corpora: [{ corpus: "wiki", outcome: "unavailable", error: "wiki unavailable" }], + warning: "Wiki corpus unavailable: wiki unavailable", + error: "wiki unavailable", + }, + }, + ])("reports a wiki-only $name", async ({ wiki, expected }) => { + prepareCase(memoryMiss, wiki); + const result = await createMemoryGetToolOrThrow().execute("call_get_wiki", { + path: lookup, + corpus: "wiki", + }); + expect(result.details).toEqual(expected); + }); + + it("uses deterministic surviving wiki content when another supplement fails", async () => { + registerMemoryCorpusSupplement("z-wiki", { + search: async () => [], + get: async () => ({ ...wikiHit, content: "Zeta entry" }), + }); + registerMemoryCorpusSupplement("m-broken", { + search: async () => [], + get: async () => { + throw new Error("broken wiki"); + }, + }); + registerMemoryCorpusSupplement("a-wiki", { search: async () => [], get: async () => wikiHit }); + + const result = await createMemoryGetToolOrThrow().execute("call_get_wiki_partial", { + path: lookup, + corpus: "wiki", + }); + + expect(result.details).toEqual({ + ...wikiPayload, + corpora: [{ corpus: "wiki", outcome: "unavailable", error: "broken wiki" }], + warning: "Wiki corpus unavailable: broken wiki", + error: "broken wiki", + }); + }); + + it.each(["wiki", "all"] as const)( + "forwards effective agent context to corpus=%s supplements", + async (corpus) => { + const get = vi.fn(async () => wikiHit); + registerMemoryCorpusSupplement("memory-wiki", { search: async () => [], get }); + const config = asOpenClawConfig({ + agents: { list: [{ id: "marketing-agent", default: true }] }, + }); + const tool = createMemoryGetTool({ + config, + agentId: " Marketing Agent ", + agentSessionKey: "agent:marketing-agent:main", + sandboxed: true, + }); + if (!tool) { + throw new Error("expected memory_get tool"); + } + + await tool.execute(`call_get_${corpus}`, { path: lookup, from: 2, lines: 4, corpus }); + + expect(get).toHaveBeenCalledWith({ + lookup, + fromLine: 2, + lineCount: 4, + agentId: "marketing-agent", + agentSessionKey: "agent:marketing-agent:main", + sandboxed: true, + }); + }, + ); + + it("settles a hanging supplement at the shared deadline and keeps memory", async () => { + vi.useFakeTimers(); + registerMemoryCorpusSupplement("memory-wiki", { + search: async () => [], + get: async () => await new Promise(() => {}), + }); + setMemoryReadFileImpl(async () => memoryHit); + + const pending = createMemoryGetToolOrThrow().execute("call_get_deadline", { + path: lookup, + corpus: "all", + }); + await vi.advanceTimersByTimeAsync(15_000); + + await expect(pending).resolves.toMatchObject({ + details: { + ...memoryHit, + corpora: [ + { corpus: "memory", outcome: "ok" }, + { + corpus: "wiki", + outcome: "unavailable", + error: "memory_get timed out after 15s", + }, + ], + }, + }); + }); + + it("cancels a hanging exact supplement read", async () => { + registerMemoryCorpusSupplement("memory-wiki", { + search: async () => [], + get: async () => await new Promise(() => {}), + }); + const controller = new AbortController(); + const pending = createMemoryGetToolOrThrow().execute( + "call_get_abort", + { path: lookup, corpus: "all" }, + controller.signal, + ); + controller.abort(new Error("cancelled")); + await expect(pending).rejects.toThrow("cancelled"); + }); +}); diff --git a/extensions/memory-core/src/memory-read-tool.ts b/extensions/memory-core/src/memory-read-tool.ts new file mode 100644 index 000000000000..495d72393cf4 --- /dev/null +++ b/extensions/memory-core/src/memory-read-tool.ts @@ -0,0 +1,95 @@ +import { formatErrorMessage } from "openclaw/plugin-sdk/error-runtime"; +import type { MemoryReadResult } from "openclaw/plugin-sdk/memory-core-host-engine-storage"; +import { jsonResult } from "openclaw/plugin-sdk/memory-core-host-runtime-core"; +import { + attemptMemoryCorpus, + composeMemoryCorpusMetadata, + readMemoryCorpusSupplements, + runMemoryCorpusDeadline, + type MemoryCorpusAttempt, +} from "./memory-corpus.js"; + +type MemoryReadRequest = { + requestedCorpus?: "memory" | "wiki" | "all"; + relPath: string; + from?: number; + lines?: number; + agentId?: string; + agentSessionKey?: string; + sandboxed?: boolean; + signal?: AbortSignal; +}; + +function readWiki(params: MemoryReadRequest, signal: AbortSignal) { + return readMemoryCorpusSupplements({ + lookup: params.relPath, + fromLine: params.from, + lineCount: params.lines, + agentId: params.agentId, + agentSessionKey: params.agentSessionKey, + sandboxed: params.sandboxed, + signal, + }); +} + +function attemptValue(attempt: MemoryCorpusAttempt): T | null { + return attempt.outcome === "not-registered" ? null : attempt.value; +} + +export async function executeWikiMemoryReadResult(params: MemoryReadRequest) { + return await runMemoryCorpusDeadline({ + operation: "memory_get", + parentSignal: params.signal, + run: async (signal) => { + const wiki = await readWiki(params, signal); + const result = + attemptValue(wiki) ?? + (wiki.outcome === "ok" + ? { status: "not_found" as const, path: params.relPath, text: "" as const } + : { path: params.relPath, text: "" }); + return jsonResult({ ...result, ...composeMemoryCorpusMetadata([wiki]) }); + }, + }); +} + +export async function executeMemoryReadResult( + params: MemoryReadRequest & { read: () => Promise }, +) { + if (params.requestedCorpus !== "all") { + try { + return jsonResult(await params.read()); + } catch (error) { + return jsonResult({ + path: params.relPath, + text: "", + disabled: true, + error: formatErrorMessage(error), + }); + } + } + return await runMemoryCorpusDeadline({ + operation: "memory_get", + parentSignal: params.signal, + run: async (signal) => { + const [memory, wiki] = await Promise.all([ + attemptMemoryCorpus({ + corpus: "memory", + signal, + unavailableValue: null, + run: params.read, + }), + readWiki(params, signal), + ]); + const memoryResult = attemptValue(memory); + const wikiResult = attemptValue(wiki); + const result = + memoryResult?.status !== "not_found" && memoryResult !== null + ? memoryResult + : (wikiResult ?? + (memory.outcome === "ok" || wiki.outcome === "ok" + ? { status: "not_found" as const, path: params.relPath, text: "" as const } + : { path: params.relPath, text: "", disabled: true })); + return jsonResult({ ...result, ...composeMemoryCorpusMetadata([memory, wiki]) }); + }, + }); +} diff --git a/extensions/memory-core/src/memory-search-tool-query.ts b/extensions/memory-core/src/memory-search-tool-query.ts index 91bb4b9a4e9c..8651146a468e 100644 --- a/extensions/memory-core/src/memory-search-tool-query.ts +++ b/extensions/memory-core/src/memory-search-tool-query.ts @@ -60,12 +60,12 @@ export async function executeMemorySearchToolQuery(params: { refreshManager: () => Promise; query: MemorySearchToolQuery; visibility: MemorySearchToolVisibility; - runWithDeadline: (task: (signal: AbortSignal) => Promise) => Promise; + signal: AbortSignal; }) { const startedAt = Date.now(); const runtimeDebug: MemorySearchRuntimeDebug[] = []; let active = params.initialManager; - const { query, runWithDeadline, visibility } = params; + const { query, signal, visibility } = params; // Product recall may index transcripts without adding them to ordinary model search. // Explicit corpus selection is authorized by the tool owner before this point. const searchSources = @@ -93,18 +93,15 @@ export async function executeMemorySearchToolQuery(params: { const searchWindow = searchesSessions ? Math.min(MEMORY_SEARCH_POST_FILTER_MAX_CANDIDATES, availableCandidates) : query.resultLimit; - const candidates = await runWithDeadline( - async (signal) => - await active.manager.search(query.text, { - maxResults: searchWindow, - minScore: query.minScore, - sessionKey: query.sessionKey, - activeProjectKeys: query.activeProjectKeys ? [...query.activeProjectKeys] : undefined, - signal, - onDebug: (debug) => runtimeDebug.push(debug), - ...(searchSources ? { sources: searchSources } : {}), - }), - ); + const candidates = await active.manager.search(query.text, { + maxResults: searchWindow, + minScore: query.minScore, + sessionKey: query.sessionKey, + activeProjectKeys: query.activeProjectKeys ? [...query.activeProjectKeys] : undefined, + signal, + onDebug: (debug) => runtimeDebug.push(debug), + ...(searchSources ? { sources: searchSources } : {}), + }); return { candidates, searchWindow }; }; @@ -141,17 +138,14 @@ export async function executeMemorySearchToolQuery(params: { }; } - let filtered = await runWithDeadline( - async () => - await filterMemorySearchHitsBySessionVisibility({ - cfg: visibility.cfg, - agentId: visibility.agentId, - requesterSessionKey: query.sessionKey, - sandboxed: visibility.sandboxed, - hits: searched.candidates, - conversationRecall: query.conversationRecall, - }), - ); + let filtered = await filterMemorySearchHitsBySessionVisibility({ + cfg: visibility.cfg, + agentId: visibility.agentId, + requesterSessionKey: query.sessionKey, + sandboxed: visibility.sandboxed, + hits: searched.candidates, + conversationRecall: query.conversationRecall, + }); if (searchSources) { const allowedSources = new Set(searchSources); filtered = filtered.filter((hit) => allowedSources.has(hit.source)); diff --git a/extensions/memory-core/src/tools.citations.test.ts b/extensions/memory-core/src/tools.citations.test.ts index d418272027db..8d1ccb8b405a 100644 --- a/extensions/memory-core/src/tools.citations.test.ts +++ b/extensions/memory-core/src/tools.citations.test.ts @@ -207,18 +207,18 @@ describe("memory tools", () => { }); }); - it("returns an explicit not-found result when the file does not exist", async () => { + it("returns an explicit not-found outcome when the file does not exist", async () => { setMemoryReadFileImpl(async (_params: MemoryReadParams) => { - return { status: "not_found", text: "", path: "memory/2026-02-19.md" }; + return { text: "", path: "memory/2026-02-19.md", status: "not_found" }; }); const tool = createMemoryGetToolOrThrow(); const result = await tool.execute("call_enoent", { path: "memory/2026-02-19.md" }); expect(result.details).toEqual({ - status: "not_found", text: "", path: "memory/2026-02-19.md", + status: "not_found", }); }); @@ -396,6 +396,7 @@ describe("memory tools", () => { snippet: "Alpha wiki entry", }, ], + corpora: [{ corpus: "wiki", outcome: "ok" }], citations: "auto", debug: undefined, fallback: undefined, @@ -446,7 +447,6 @@ describe("memory tools", () => { agentId: "marketing-agent", agentSessionKey: "agent:marketing-agent:main", sandboxed: true, - corpus, }); }, ); @@ -639,10 +639,17 @@ describe("memory tools", () => { }); await vi.advanceTimersByTimeAsync(15_000); const stalledAllResult = await stalledAllResultPromise; - expectUnavailableMemorySearchDetails(stalledAllResult.details, { - error: "memory_search timed out after 15s", - warning: "Memory search is unavailable due to an embedding/provider error.", - action: "Check embedding provider configuration and retry memory_search.", + expect(stalledAllResult.details).toMatchObject({ + results: [{ corpus: "memory", path: "MEMORY.md" }], + corpora: [ + { corpus: "memory", outcome: "ok" }, + { + corpus: "wiki", + outcome: "unavailable", + error: "memory_search timed out after 15s", + }, + ], + warning: expect.stringContaining("Wiki corpus unavailable"), }); const memoryResult = await tool.execute("call_memory_after_stalled_wiki", { @@ -658,6 +665,23 @@ describe("memory tools", () => { } }); + it("records an unregistered requested wiki corpus without hiding memory results", async () => { + const tool = createMemorySearchToolOrThrow(); + const result = await tool.execute("call_all_without_wiki", { + query: "alpha", + corpus: "all", + }); + + expect(result.details).toMatchObject({ + results: [{ corpus: "memory", path: "MEMORY.md" }], + corpora: [ + { corpus: "memory", outcome: "ok" }, + { corpus: "wiki", outcome: "not-registered" }, + ], + warning: expect.stringContaining("Wiki corpus is not registered"), + }); + }); + it("surfaces a memory-corpus warning when corpus=all hits a returned manager error", async () => { setMemorySearchManagerImpl(async () => ({ error: "sqlite support missing" })); registerMemoryCorpusSupplement("memory-wiki", { @@ -716,10 +740,17 @@ describe("memory tools", () => { }); await vi.advanceTimersByTimeAsync(15_000); const stalledAllResult = await stalledAllResultPromise; - expectUnavailableMemorySearchDetails(stalledAllResult.details, { - error: "memory_search timed out after 15s", - warning: "Memory search is unavailable due to an embedding/provider error.", - action: "Check embedding provider configuration and retry memory_search.", + expect(stalledAllResult.details).toMatchObject({ + results: [{ corpus: "wiki", path: "entities/alpha.md" }], + corpora: [ + { + corpus: "memory", + outcome: "unavailable", + error: "memory_search timed out after 15s", + }, + { corpus: "wiki", outcome: "ok" }, + ], + warning: expect.stringContaining("Memory corpus unavailable"), }); const wikiOnlyResult = await tool.execute("call_all_after_stalled_memory", { @@ -728,203 +759,25 @@ describe("memory tools", () => { }); const details = wikiOnlyResult.details as { results: Array<{ corpus: string; path: string }>; + corpora: Array<{ corpus: string; outcome: string; error?: string }>; + warning?: string; }; expect(details.results.map((entry) => [entry.corpus, entry.path])).toEqual([ ["wiki", "entities/alpha.md"], ]); + expect(details.corpora).toEqual([ + { + corpus: "memory", + outcome: "unavailable", + error: "memory_search timed out after 15s", + }, + { corpus: "wiki", outcome: "ok" }, + ]); + expect(details.warning).toContain("Memory corpus unavailable"); + expect(details.warning).toContain("memory_search timed out after 15s"); expect(searchCalls).toBe(1); } finally { vi.useRealTimers(); } }); - - it("falls back to a wiki corpus supplement for memory_get corpus=all", async () => { - setMemoryReadFileImpl(async () => { - throw new Error("path required"); - }); - registerMemoryCorpusSupplement("memory-wiki", { - search: async () => [], - get: async () => ({ - corpus: "wiki", - path: "entities/alpha.md", - title: "Alpha", - kind: "entity", - content: "Alpha wiki entry", - fromLine: 3, - lineCount: 5, - }), - }); - - const tool = createMemoryGetToolOrThrow(); - const result = await tool.execute("call_get_all_fallback", { - path: "entities/alpha.md", - from: 3, - lines: 5, - corpus: "all", - }); - - expect(result.details).toEqual({ - corpus: "wiki", - path: "entities/alpha.md", - title: "Alpha", - kind: "entity", - text: "Alpha wiki entry", - fromLine: 3, - lineCount: 5, - }); - }); - - it.each(["wiki", "all"] as const)( - "forwards effective agent context to memory_get corpus=%s supplements", - async (corpus) => { - if (corpus === "all") { - setMemoryReadFileImpl(async () => { - throw new Error("memory path missing"); - }); - } - const get = vi.fn(async () => ({ - corpus: "wiki" as const, - path: "entities/alpha.md", - content: "Alpha wiki entry", - fromLine: 2, - lineCount: 4, - })); - registerMemoryCorpusSupplement("memory-wiki", { - search: async () => [], - get, - }); - const config = asOpenClawConfig({ - agents: { list: [{ id: "marketing-agent", default: true }] }, - }); - const tool = createMemoryGetTool({ - config, - agentId: " Marketing Agent ", - agentSessionKey: "agent:marketing-agent:main", - sandboxed: true, - }); - if (!tool) { - throw new Error("expected memory_get tool"); - } - - await tool.execute(`call_get_${corpus}`, { - path: "entities/alpha.md", - from: 2, - lines: 4, - corpus, - }); - - expect(get).toHaveBeenCalledWith({ - lookup: "entities/alpha.md", - fromLine: 2, - lineCount: 4, - agentId: "marketing-agent", - agentSessionKey: "agent:marketing-agent:main", - sandboxed: true, - corpus, - }); - }, - ); - - it("falls back to a wiki corpus supplement when memory_get corpus=all misses memory without throwing", async () => { - setMemoryReadFileImpl(async (params: MemoryReadParams) => ({ - status: "not_found", - text: "", - path: params.relPath, - })); - registerMemoryCorpusSupplement("memory-wiki", { - search: async () => [], - get: async () => ({ - corpus: "wiki", - path: "memory/entities/alpha.md", - title: "Alpha", - kind: "entity", - content: "Alpha wiki entry after empty miss", - fromLine: 3, - lineCount: 5, - }), - }); - - const tool = createMemoryGetToolOrThrow(); - const result = await tool.execute("call_get_all_empty_miss_fallback", { - path: "memory/entities/alpha.md", - from: 3, - lines: 5, - corpus: "all", - }); - - expect(result.details).toEqual({ - corpus: "wiki", - path: "memory/entities/alpha.md", - title: "Alpha", - kind: "entity", - text: "Alpha wiki entry after empty miss", - fromLine: 3, - lineCount: 5, - }); - }); - - it("preserves an empty in-file range for memory_get corpus=all", async () => { - setMemoryReadFileImpl(async (params: MemoryReadParams) => ({ - status: "ok", - text: "", - path: params.relPath, - from: params.from ?? 1, - lines: 0, - })); - const getSupplement = vi.fn(async () => ({ - corpus: "wiki" as const, - path: "memory/entities/alpha.md", - title: "Alpha", - kind: "entity", - content: "Alpha wiki entry", - fromLine: 10, - lineCount: 5, - })); - registerMemoryCorpusSupplement("memory-wiki", { - search: async () => [], - get: getSupplement, - }); - - const tool = createMemoryGetToolOrThrow(); - const result = await tool.execute("call_get_all_empty_range", { - path: "memory/entities/alpha.md", - from: 10, - lines: 5, - corpus: "all", - }); - - expect(result.details).toEqual({ - status: "ok", - text: "", - path: "memory/entities/alpha.md", - from: 10, - lines: 0, - }); - expect(getSupplement).not.toHaveBeenCalled(); - }); - - it("returns the primary error when a corpus=all supplement fallback throws", async () => { - setMemoryReadFileImpl(async () => { - throw new Error("primary read failed"); - }); - registerMemoryCorpusSupplement("memory-wiki", { - search: async () => [], - get: async () => { - throw new Error("supplement lookup failed"); - }, - }); - - const tool = createMemoryGetToolOrThrow(); - const result = await tool.execute("call_get_all_supplement_throws", { - path: "entities/alpha.md", - corpus: "all", - }); - - expect(result.details).toEqual({ - path: "entities/alpha.md", - text: "", - disabled: true, - error: "primary read failed", - }); - }); }); diff --git a/extensions/memory-core/src/tools.shared.ts b/extensions/memory-core/src/tools.shared.ts index b0b03460e5e8..102c740b4880 100644 --- a/extensions/memory-core/src/tools.shared.ts +++ b/extensions/memory-core/src/tools.shared.ts @@ -2,10 +2,8 @@ import { optionalFiniteNumberSchema, stringEnum } from "openclaw/plugin-sdk/channel-actions"; import { createLazyRuntimeModule } from "openclaw/plugin-sdk/lazy-runtime"; import { - listMemoryCorpusSupplements, resolveMemorySearchConfig, resolveSessionAgentIds, - type MemoryCorpusSearchResult, type AnyAgentTool, type OpenClawConfig, } from "openclaw/plugin-sdk/memory-core-host-runtime-core"; @@ -51,10 +49,7 @@ function resolveMemoryToolContext(options: MemoryToolOptions) { config: cfg, agentId: options.agentId, }); - if (!resolveMemorySearchConfig(cfg, agentId)) { - return null; - } - return { cfg, agentId }; + return resolveMemorySearchConfig(cfg, agentId) ? { cfg, agentId } : null; } export async function getMemoryManagerContextWithPurpose(params: { @@ -163,54 +158,3 @@ export function buildMemorySearchUnavailableResult( }, }; } - -export async function searchMemoryCorpusSupplements(params: { - query: string; - maxResults?: number; - agentId?: string; - agentSessionKey?: string; - sandboxed?: boolean; - corpus?: "memory" | "wiki" | "all" | "sessions"; -}): Promise { - if (params.corpus === "memory" || params.corpus === "sessions") { - return []; - } - const supplements = listMemoryCorpusSupplements(); - if (supplements.length === 0) { - return []; - } - const results = ( - await Promise.all( - supplements.map(async (registration) => await registration.supplement.search(params)), - ) - ).flat(); - return results - .toSorted((left, right) => { - if (left.score !== right.score) { - return right.score - left.score; - } - return left.path.localeCompare(right.path); - }) - .slice(0, Math.max(1, params.maxResults ?? 10)); -} - -export async function getMemoryCorpusSupplementResult(params: { - lookup: string; - fromLine?: number; - lineCount?: number; - agentId?: string; - agentSessionKey?: string; - sandboxed?: boolean; - corpus?: "memory" | "wiki" | "all" | "sessions"; -}) { - if (params.corpus === "memory" || params.corpus === "sessions") { - return null; - } - for (const registration of listMemoryCorpusSupplements()) { - const result = await registration.supplement.get(params); - if (result) { - return result; - } - } - return null; -} diff --git a/extensions/memory-core/src/tools.ts b/extensions/memory-core/src/tools.ts index 9bbb9422bcc5..614f2f0b8bbb 100644 --- a/extensions/memory-core/src/tools.ts +++ b/extensions/memory-core/src/tools.ts @@ -3,7 +3,6 @@ import { formatErrorMessage } from "openclaw/plugin-sdk/error-runtime"; import { resolveMemorySearchStaleness, stripMemoryAnnotationCarriers, - type MemoryReadResult, type MemorySource, } from "openclaw/plugin-sdk/memory-core-host-engine-storage"; import { @@ -23,6 +22,15 @@ import { resolveMemoryDeepDreamingConfig, } from "openclaw/plugin-sdk/memory-core-host-status"; import type { OpenClawPluginToolContext } from "openclaw/plugin-sdk/plugin-entry"; +import { + attemptMemoryCorpus, + composeMemoryCorpusMetadata, + runMemoryCorpusDeadline, + searchMemoryCorpusSupplements, + unavailableMemoryCorpus, + type MemoryCorpusAttempt, +} from "./memory-corpus.js"; +import { executeMemoryReadResult, executeWikiMemoryReadResult } from "./memory-read-tool.js"; import { buildPausedMemoryIndexUnavailableResult, executeMemorySearchToolQuery, @@ -42,12 +50,10 @@ import { import { buildMemorySearchUnavailableResult, createMemoryTool, - getMemoryCorpusSupplementResult, getMemoryManagerContextWithPurpose, loadMemoryToolRuntime, MemoryGetSchema, MemorySearchSchema, - searchMemoryCorpusSupplements, } from "./tools.shared.js"; type MemorySearchToolResult = @@ -58,6 +64,17 @@ type ActiveMemoryManagerContext = Extract>["debug"] >; +type PrimaryMemorySearchValue = { + results: Array; + rawResults: MemorySearchResult[]; + provider?: string; + model?: string; + fallback?: unknown; + mode?: string; + staleness?: Exclude, null>; + debug?: MemorySearchToolQueryDebug & { toolMs?: number; outsideSearchMs?: number }; + unavailableResult?: ReturnType; +}; const MEMORY_SEARCH_TOOL_COOLDOWN_MS = 60_000; @@ -137,7 +154,10 @@ async function closeMemoryManagers( await Promise.allSettled(pending); }, }); - } catch { + } catch (error) { + if (parentSignal?.aborted) { + throw error; + } // Search results should not be hidden by best-effort transient cleanup. } } @@ -249,108 +269,6 @@ function queueShortTermRecallTracking(params: { }); } -async function getSupplementMemoryReadResult(params: { - relPath: string; - from?: number; - lines?: number; - agentId?: string; - agentSessionKey?: string; - sandboxed?: boolean; - corpus?: "memory" | "wiki" | "all"; -}) { - const supplement = await getMemoryCorpusSupplementResult({ - lookup: params.relPath, - fromLine: params.from, - lineCount: params.lines, - agentId: params.agentId, - agentSessionKey: params.agentSessionKey, - sandboxed: params.sandboxed, - corpus: params.corpus, - }); - if (!supplement) { - return null; - } - const { content, ...rest } = supplement; - return { - ...rest, - text: content, - }; -} - -async function resolveMemoryReadFailureResult(params: { - error: unknown; - requestedCorpus?: "memory" | "wiki" | "all"; - relPath: string; - from?: number; - lines?: number; - agentId?: string; - agentSessionKey?: string; - sandboxed?: boolean; -}) { - if (params.requestedCorpus === "all") { - try { - const supplement = await getSupplementMemoryReadResult({ - relPath: params.relPath, - from: params.from, - lines: params.lines, - agentId: params.agentId, - agentSessionKey: params.agentSessionKey, - sandboxed: params.sandboxed, - corpus: params.requestedCorpus, - }); - if (supplement) { - return jsonResult(supplement); - } - } catch { - // Supplement lookup is best-effort after the primary memory read failed. - // Preserve the original structured error instead of rejecting the tool call. - } - } - const message = formatErrorMessage(params.error); - return jsonResult({ path: params.relPath, text: "", disabled: true, error: message }); -} - -async function executeMemoryReadResult(params: { - read: () => Promise; - requestedCorpus?: "memory" | "wiki" | "all"; - relPath: string; - from?: number; - lines?: number; - agentId?: string; - agentSessionKey?: string; - sandboxed?: boolean; -}) { - try { - const result = await params.read(); - if (params.requestedCorpus === "all" && result.status === "not_found") { - const supplement = await getSupplementMemoryReadResult({ - relPath: params.relPath, - from: params.from, - lines: params.lines, - agentId: params.agentId, - agentSessionKey: params.agentSessionKey, - sandboxed: params.sandboxed, - corpus: params.requestedCorpus, - }); - if (supplement) { - return jsonResult(supplement); - } - } - return jsonResult(result); - } catch (error) { - return await resolveMemoryReadFailureResult({ - error, - requestedCorpus: params.requestedCorpus, - relPath: params.relPath, - from: params.from, - lines: params.lines, - agentId: params.agentId, - agentSessionKey: params.agentSessionKey, - sandboxed: params.sandboxed, - }); - } -} - export function createMemorySearchTool(options: { config?: OpenClawConfig; getConfig?: () => OpenClawConfig | undefined; @@ -367,7 +285,7 @@ export function createMemorySearchTool(options: { label: "Memory Search", name: "memory_search", description: - "Mandatory recall step: semantically search MEMORY.md + memory/*.md (and optional session transcripts) before answering questions about prior work, decisions, dates, people, preferences, or todos. Optional `corpus=wiki` or `corpus=all` also searches registered compiled-wiki supplements. `corpus=memory` restricts hits to indexed memory files (excludes session transcript chunks from ranking). `corpus=sessions` restricts hits to indexed session transcripts (same visibility rules as session history tools). If response has disabled=true or stale=true, you must tell the user and include the warning/action guidance.", + "Mandatory recall step: semantically search MEMORY.md + memory/*.md (and optional session transcripts) before answering questions about prior work, decisions, dates, people, preferences, or todos. Optional `corpus=wiki` or `corpus=all` also searches registered compiled-wiki supplements. `corpus=memory` restricts hits to indexed memory files (excludes session transcript chunks from ranking). `corpus=sessions` restricts hits to indexed session transcripts (same visibility rules as session history tools). Corpus warnings mean the returned results are partial and must be surfaced to the user. If response has disabled=true or stale=true, tell the user and include the warning/action guidance.", parameters: MemorySearchSchema, execute: ({ cfg, agentId }) => @@ -394,278 +312,231 @@ export function createMemorySearchTool(options: { }); const cooldown = requestedCorpus === "wiki" ? undefined : readMemorySearchToolCooldown(cooldownKey); - let activeUnavailablePhase: "memory" | "supplement" | undefined; - let failedUnavailablePhase: "memory" | "supplement" | undefined; - const runUnavailablePhase = async ( - phase: "memory" | "supplement", - task: () => Promise, - ): Promise => { - activeUnavailablePhase = phase; - try { - return await task(); - } catch (error) { - failedUnavailablePhase = phase; - throw error; - } finally { - if (activeUnavailablePhase === phase) { - activeUnavailablePhase = undefined; + const toolStartedAt = Date.now(); + const searchesMemory = requestedCorpus !== "wiki"; + const searchesWiki = requestedCorpus === "wiki" || requestedCorpus === "all"; + const memoryManagerPurpose = options.oneShotCliRun ? "cli" : undefined; + const memoryManagersToClose = new Set(); + let cleanupStarted = false; + const trackMemoryManager = (context: MemoryManagerContext): MemoryManagerContext => { + if (memoryManagerPurpose === "cli" && isActiveMemoryManagerContext(context)) { + if (cleanupStarted) { + void closeMemoryManagers([context.manager]); + } else { + memoryManagersToClose.add(context.manager); } } + return context; }; - const runWithDefaultDeadline = async ( - task: (signal: AbortSignal) => Promise, - ): Promise => - await runMemorySearchWithDeadline({ - timeoutMs: DEFAULT_MEMORY_SEARCH_TIMEOUT_MS, - parentSignal: callerSignal, - run: task, - }); - const runMemorySearchTool = async () => { - const toolStartedAt = Date.now(); - const shouldQuerySupplements = requestedCorpus === "wiki" || requestedCorpus === "all"; - const shouldQueryMemory = requestedCorpus !== "wiki" && !cooldown; - if (cooldown && !shouldQuerySupplements) { - return jsonResult(buildMemorySearchUnavailableResult(cooldown.error)); + const searchMemory = async ( + signal: AbortSignal, + ): Promise> => { + if (cooldown) { + return unavailableMemoryCorpus("memory", null, cooldown.error); } - const memoryManagerPurpose = options.oneShotCliRun ? "cli" : undefined; - const memoryManagersToClose = new Set(); - let cleanupStarted = false; - const trackMemoryManager = (context: MemoryManagerContext): MemoryManagerContext => { - if (memoryManagerPurpose === "cli" && isActiveMemoryManagerContext(context)) { - if (cleanupStarted) { - // Setup can settle after its deadline. Close that late transient - // manager instead of leaking it after the tool has returned. - void closeMemoryManagers([context.manager]); - } else { - memoryManagersToClose.add(context.manager); - } - } - return context; - }; - try { - const memorySetup = shouldQueryMemory - ? await runUnavailablePhase( - "memory", - async () => - await runWithDefaultDeadline(async () => { - const context = trackMemoryManager( - await getMemoryManagerContextWithPurpose({ - cfg, - agentId, - purpose: memoryManagerPurpose, - acquireLocalService: options.acquireLocalService, - }), - ); - return { context }; - }), - ) - : null; - const memory = memorySetup?.context ?? null; - let memoryCorpusUnavailable: string | undefined; - if (shouldQueryMemory && memory && "error" in memory) { - recordMemorySearchToolCooldown( - cooldownKey, - memory.error ?? "memory search unavailable", + const attempted = await attemptMemoryCorpus + > | null>({ + corpus: "memory", + signal, + unavailableValue: null, + run: async () => { + const memory = trackMemoryManager( + await getMemoryManagerContextWithPurpose({ + cfg, + agentId, + purpose: memoryManagerPurpose, + acquireLocalService: options.acquireLocalService, + }), ); - if (!shouldQuerySupplements) { - return jsonResult(buildMemorySearchUnavailableResult(memory.error)); + if ("error" in memory) { + throw new Error(memory.error ?? "memory search unavailable"); } - // corpus=all still serves wiki supplements, but the omitted memory - // corpus must be recorded or the degraded search reads as complete. - memoryCorpusUnavailable = memory.error ?? "memory search unavailable"; - } - - const citationsMode = resolveMemoryCitationsMode(cfg); - const includeCitations = shouldIncludeCitations({ - mode: citationsMode, - sessionKey: options.agentSessionKey, - }); - const pluginConfig = resolveMemoryDreamingPluginConfig(cfg); - const dreamingEnabled = resolveMemoryDreamingConfig({ - pluginConfig, - cfg, - }).enabled; - const dreaming = resolveMemoryDeepDreamingConfig({ - pluginConfig, - cfg, - }); - let rawResults: MemorySearchResult[] = []; - let surfacedMemoryResults: Array = []; - let provider: string | undefined, model: string | undefined; - let fallback: unknown; - let searchMode: string | undefined, pausedIndexIdentityReason: string | undefined; - let staleness: - | Exclude, null> - | undefined; - let searchDebug: - | (MemorySearchToolQueryDebug & { toolMs?: number; outsideSearchMs?: number }) - | undefined; - if (shouldQueryMemory && memorySetup && memory && !("error" in memory)) { - await runUnavailablePhase("memory", async () => { - const memorySearchConfig = resolveMemorySearchConfig(cfg, agentId); - const defaultSearchSources = memorySearchConfig?.searchSources; - const explicitSearchSources: MemorySource[] | undefined = - requestedCorpus === "sessions" && - (options.conversationRecall || defaultSearchSources?.includes("sessions")) - ? (["sessions"] as MemorySource[]) - : requestedCorpus === "memory" - ? (["memory"] as MemorySource[]) - : undefined; - const resultLimit = maxResults ?? memorySearchConfig?.query.maxResults ?? 10; - const executed = await executeMemorySearchToolQuery({ - initialManager: { - manager: memory.manager, - managerMs: memory.debug?.managerMs, - }, - refreshManager: async () => { - const refreshed = await runWithDefaultDeadline(async () => - trackMemoryManager( - await getMemoryManagerContextWithPurpose({ - cfg, - agentId, - purpose: memoryManagerPurpose, - acquireLocalService: options.acquireLocalService, - }), - ), - ); - if ("error" in refreshed) { - return null; - } - return { - manager: refreshed.manager, - managerMs: refreshed.debug?.managerMs, - }; - }, - query: { - text: query, - resultLimit, - minScore, - explicitSources: explicitSearchSources, - defaultSources: defaultSearchSources, - indexedSources: memorySearchConfig?.sources, - requestedCorpus, - sessionKey: options.agentSessionKey, - activeProjectKeys: options.activeProjectKeys, - conversationRecall: options.conversationRecall, - }, - visibility: { - cfg, - agentId, - sandboxed: options.sandboxed === true, - }, - runWithDeadline: runWithDefaultDeadline, - }); - pausedIndexIdentityReason = executed.pausedIndexIdentityReason; - if (pausedIndexIdentityReason) { - return; - } - rawResults = executed.rawResults; - const status = executed.status; - staleness = resolveMemorySearchStaleness(status, agentId) ?? undefined; - const payloadResults = rawResults.map((result) => ({ - ...result, - snippet: stripMemoryAnnotationCarriers(result.snippet), - })); - const decorated = decorateCitations(payloadResults, includeCitations); - const memoryResults = decorated; - surfacedMemoryResults = memoryResults.map((result) => ({ - ...result, - corpus: result.source, - })); - if (dreamingEnabled) { - queueShortTermRecallTracking({ - workspaceDir: status.workspaceDir, - query, - rawResults, - surfacedResults: memoryResults, - timezone: dreaming.timezone, - }); - } - provider = status.provider; - model = status.model; - fallback = status.fallback; - searchMode = executed.searchMode; - searchDebug = executed.debug; + const settings = resolveMemorySearchConfig(cfg, agentId); + const defaultSources = settings?.searchSources; + const explicitSources: MemorySource[] | undefined = + requestedCorpus === "sessions" && + (options.conversationRecall || defaultSources?.includes("sessions")) + ? ["sessions"] + : requestedCorpus === "memory" + ? ["memory"] + : undefined; + return await executeMemorySearchToolQuery({ + initialManager: { manager: memory.manager, managerMs: memory.debug?.managerMs }, + refreshManager: async () => { + const refreshed = trackMemoryManager( + await getMemoryManagerContextWithPurpose({ + cfg, + agentId, + purpose: memoryManagerPurpose, + acquireLocalService: options.acquireLocalService, + }), + ); + return "error" in refreshed + ? null + : { manager: refreshed.manager, managerMs: refreshed.debug?.managerMs }; + }, + query: { + text: query, + resultLimit: maxResults ?? settings?.query.maxResults ?? 10, + minScore, + explicitSources, + defaultSources, + indexedSources: settings?.sources, + requestedCorpus, + sessionKey: options.agentSessionKey, + activeProjectKeys: options.activeProjectKeys, + conversationRecall: options.conversationRecall, + }, + visibility: { cfg, agentId, sandboxed: options.sandboxed === true }, + signal, }); - if (pausedIndexIdentityReason) { - return jsonResult( - buildPausedMemoryIndexUnavailableResult(pausedIndexIdentityReason), - ); - } + }, + }); + if (attempted.outcome !== "ok") { + if (callerSignal?.aborted) { + throw resolveMemorySearchAbortError(callerSignal); } - const supplementResults = shouldQuerySupplements - ? await runUnavailablePhase( - "supplement", - async () => - await runWithDefaultDeadline( - async () => - await searchMemoryCorpusSupplements({ - query, - maxResults, - agentId, - agentSessionKey: options.agentSessionKey, - sandboxed: options.sandboxed, - corpus: requestedCorpus, - }), - ), - ) - : []; - // 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 results = mergeMemorySearchCorpusResults({ - memoryResults: surfacedMemoryResults, - supplementResults, - maxResults: effectiveMax, - balanceCorpora: requestedCorpus === "all", - }); - if (searchDebug) { - const finalToolMs = Math.max(0, Date.now() - toolStartedAt); - searchDebug = { - ...searchDebug, - toolMs: finalToolMs, - outsideSearchMs: Math.max(0, finalToolMs - searchDebug.searchMs), - }; - } - return jsonResult({ - results, - provider, - model, - fallback, - citations: citationsMode, - mode: searchMode, - ...(memoryCorpusUnavailable - ? { - warning: `Memory corpus unavailable; results cover wiki supplements only: ${memoryCorpusUnavailable}`, - } - : {}), - ...staleness, - debug: searchDebug, - }); - } finally { - cleanupStarted = true; - await closeMemoryManagers(memoryManagersToClose, callerSignal); + const error = + attempted.outcome === "unavailable" ? attempted.error : "memory search unavailable"; + recordMemorySearchToolCooldown(cooldownKey, error); + return unavailableMemoryCorpus("memory", null, error); } + const executed = attempted.value!; + if (executed.pausedIndexIdentityReason) { + const reason = executed.pausedIndexIdentityReason; + return unavailableMemoryCorpus( + "memory", + { + results: [], + rawResults: [], + unavailableResult: buildPausedMemoryIndexUnavailableResult(reason), + }, + reason, + ); + } + const citationsMode = resolveMemoryCitationsMode(cfg); + const includeCitations = shouldIncludeCitations({ + mode: citationsMode, + sessionKey: options.agentSessionKey, + }); + const rawResults = executed.rawResults; + const memoryResults = decorateCitations( + rawResults.map((result) => ({ + ...result, + snippet: stripMemoryAnnotationCarriers(result.snippet), + })), + includeCitations, + ); + const status = executed.status; + if ( + resolveMemoryDreamingConfig({ + pluginConfig: resolveMemoryDreamingPluginConfig(cfg), + cfg, + }).enabled + ) { + queueShortTermRecallTracking({ + workspaceDir: status.workspaceDir, + query, + rawResults, + surfacedResults: memoryResults, + timezone: resolveMemoryDeepDreamingConfig({ + pluginConfig: resolveMemoryDreamingPluginConfig(cfg), + cfg, + }).timezone, + }); + } + return { + corpus: "memory", + outcome: "ok", + value: { + results: memoryResults.map((result) => + Object.assign(result, { corpus: result.source }), + ), + rawResults, + provider: status.provider, + model: status.model, + fallback: status.fallback, + mode: executed.searchMode, + staleness: resolveMemorySearchStaleness(status, agentId) ?? undefined, + debug: executed.debug, + }, + }; }; try { - const result = await runMemorySearchTool(); - if (callerSignal?.aborted) { - throw resolveMemorySearchAbortError(callerSignal); - } - return result; + return await runMemoryCorpusDeadline({ + operation: "memory_search", + parentSignal: callerSignal, + run: async (signal) => { + const [memory, wiki] = await Promise.all([ + searchesMemory ? searchMemory(signal) : Promise.resolve(null), + searchesWiki + ? searchMemoryCorpusSupplements({ + query, + maxResults, + agentId, + agentSessionKey: options.agentSessionKey, + sandboxed: options.sandboxed, + signal, + }) + : Promise.resolve(null), + ]); + const memoryValue = memory?.outcome === "not-registered" ? null : memory?.value; + if (searchesMemory && !searchesWiki && memory?.outcome === "unavailable") { + return jsonResult( + memoryValue?.unavailableResult ?? + buildMemorySearchUnavailableResult(memory.error), + ); + } + const wikiResults = wiki?.outcome === "not-registered" ? [] : (wiki?.value ?? []); + const results = mergeMemorySearchCorpusResults({ + memoryResults: memoryValue?.results ?? [], + supplementResults: wikiResults, + maxResults: Math.max(1, maxResults ?? 10), + balanceCorpora: requestedCorpus === "all", + }); + const attempts = [ + ...(requestedCorpus === "all" && memory ? [memory] : []), + ...(wiki ? [wiki] : []), + ]; + const staleness = memoryValue?.staleness; + const metadata = composeMemoryCorpusMetadata( + attempts, + staleness?.warning ? [staleness.warning] : [], + ); + const elapsed = Math.max(0, Date.now() - toolStartedAt); + const debug = memoryValue?.debug + ? { + ...memoryValue.debug, + toolMs: elapsed, + outsideSearchMs: Math.max(0, elapsed - memoryValue.debug.searchMs), + } + : undefined; + return jsonResult({ + results, + provider: memoryValue?.provider, + model: memoryValue?.model, + fallback: memoryValue?.fallback, + citations: resolveMemoryCitationsMode(cfg), + mode: memoryValue?.mode, + ...(attempts.length > 0 ? metadata : {}), + ...staleness, + debug, + }); + }, + }); } catch (error) { if (callerSignal?.aborted) { throw resolveMemorySearchAbortError(callerSignal); } - const unavailablePhase = failedUnavailablePhase ?? activeUnavailablePhase; - const shouldRecordCooldown = - requestedCorpus !== "wiki" && - (requestedCorpus !== "all" || unavailablePhase === "memory"); const message = formatErrorMessage(error); - if (shouldRecordCooldown) { + if (requestedCorpus !== "wiki") { recordMemorySearchToolCooldown(cooldownKey, message); } return jsonResult(buildMemorySearchUnavailableResult(message)); + } finally { + cleanupStarted = true; + await closeMemoryManagers(memoryManagersToClose, callerSignal); } }, }); @@ -684,11 +555,11 @@ export function createMemoryGetTool(options: { label: "Memory Get", name: "memory_get", description: - "Safe exact excerpt read from MEMORY.md or memory/*.md. Defaults to a bounded excerpt when lines are omitted, includes truncation/continuation info when more content exists, and `corpus=wiki` reads from registered compiled-wiki supplements.", + "Safe exact excerpt read from MEMORY.md or memory/*.md. Defaults to a bounded excerpt when lines are omitted, includes truncation/continuation info when more content exists, and `corpus=wiki` reads from registered compiled-wiki supplements. A response with status=not_found means every requested available corpus missed; corpus warnings mean coverage is partial and must be surfaced to the user.", parameters: MemoryGetSchema, execute: ({ cfg, agentId }) => - async (_toolCallId, params) => { + async (_toolCallId, params, callerSignal) => { const rawParams = asToolParamsRecord(params); const relPath = readStringParam(rawParams, "path", { required: true }); const from = readPositiveIntegerParam(rawParams, "from"); @@ -696,23 +567,16 @@ export function createMemoryGetTool(options: { const requestedCorpus = readCorpusParam(rawParams, ["memory", "wiki", "all"]); const { readAgentMemoryFile } = await loadMemoryToolRuntime(); if (requestedCorpus === "wiki") { - const supplement = await getSupplementMemoryReadResult({ + return await executeWikiMemoryReadResult({ relPath, from: from ?? undefined, lines: lines ?? undefined, agentId, agentSessionKey: options.agentSessionKey, sandboxed: options.sandboxed, - corpus: requestedCorpus, + requestedCorpus, + signal: callerSignal, }); - return jsonResult( - supplement ?? { - path: relPath, - text: "", - disabled: true, - error: "wiki corpus result not found", - }, - ); } return await executeMemoryReadResult({ read: async () => @@ -730,6 +594,7 @@ export function createMemoryGetTool(options: { agentId, agentSessionKey: options.agentSessionKey, sandboxed: options.sandboxed, + signal: callerSignal, }); }, }); diff --git a/extensions/memory-wiki/src/agent-vault-isolation.test.ts b/extensions/memory-wiki/src/agent-vault-isolation.test.ts index 19560b658294..ab6397ef51b1 100644 --- a/extensions/memory-wiki/src/agent-vault-isolation.test.ts +++ b/extensions/memory-wiki/src/agent-vault-isolation.test.ts @@ -246,6 +246,7 @@ describe("agent-scoped memory-wiki tools", () => { corpus: "wiki", }); expect(assertToolDetailsRecord(ownMemoryGet.details)).toMatchObject({ + status: "ok", corpus: "wiki", path: agent.pagePath, text: expect.stringContaining(agent.sentinel), @@ -258,8 +259,8 @@ describe("agent-scoped memory-wiki tools", () => { expect(assertToolDetailsRecord(foreignMemoryGet.details)).toMatchObject({ path: foreignAgent.pagePath, text: "", - disabled: true, - error: "wiki corpus result not found", + status: "not_found", + corpora: [{ corpus: "wiki", outcome: "ok" }], }); } } finally { diff --git a/extensions/memory-wiki/src/corpus-supplement.visibility.test.ts b/extensions/memory-wiki/src/corpus-supplement.visibility.test.ts index c8211f3fc2b6..8b06a77a2d9c 100644 --- a/extensions/memory-wiki/src/corpus-supplement.visibility.test.ts +++ b/extensions/memory-wiki/src/corpus-supplement.visibility.test.ts @@ -219,15 +219,16 @@ describe("memory-wiki corpus supplement visibility", () => { expect(assertToolDetailsRecord(foreignGet.details)).toMatchObject({ text: "", - disabled: true, - error: "wiki corpus result not found", + status: "not_found", + corpora: [{ corpus: "wiki", outcome: "ok" }], }); expect(assertToolDetailsRecord(unownedGet.details)).toMatchObject({ text: "", - disabled: true, - error: "wiki corpus result not found", + status: "not_found", + corpora: [{ corpus: "wiki", outcome: "ok" }], }); expect(assertToolDetailsRecord(ownGet.details)).toMatchObject({ + status: "ok", path: "sources/main-private.md", text: expect.stringContaining("REDACTED-OWN-MARKER"), }); @@ -236,6 +237,7 @@ describe("memory-wiki corpus supplement visibility", () => { expect.objectContaining({ path: "sources/main-private.md" }), ]); expect(assertToolDetailsRecord(openForeignGet.details)).toMatchObject({ + status: "ok", path: "sources/secondary-private.md", text: expect.stringContaining("REDACTED-FOREIGN-MARKER"), });