Skip to content

Commit 1407fef

Browse files
yu-xin-cyu-xin-c
committed
fix(memory-core): honor qmd search timeouts
1 parent c773d8c commit 1407fef

3 files changed

Lines changed: 162 additions & 12 deletions

File tree

‎extensions/memory-core/src/memory-tool-manager-mock.ts‎

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -68,15 +68,48 @@ const readAgentMemoryFileMock = vi.fn(
6868
async (params: MemoryReadParams) => await readFileImpl(params),
6969
);
7070

71+
function resolvePositiveIntegerConfig(raw: unknown, fallback: number): number {
72+
return typeof raw === "number" && Number.isFinite(raw) && raw > 0
73+
? Math.max(1, Math.floor(raw))
74+
: fallback;
75+
}
76+
7177
vi.mock("./tools.runtime.js", () => ({
7278
resolveMemoryBackendConfig: ({
7379
cfg,
7480
}: {
75-
cfg?: { memory?: { backend?: string; qmd?: unknown } };
76-
}) => ({
77-
backend,
78-
qmd: cfg?.memory?.qmd,
79-
}),
81+
cfg?: {
82+
memory?: {
83+
backend?: string;
84+
qmd?: {
85+
limits?: {
86+
maxResults?: unknown;
87+
maxSnippetChars?: unknown;
88+
maxInjectedChars?: unknown;
89+
timeoutMs?: unknown;
90+
};
91+
};
92+
};
93+
};
94+
}) => {
95+
const limits = cfg?.memory?.qmd?.limits;
96+
return {
97+
backend,
98+
qmd:
99+
backend === "qmd"
100+
? {
101+
...cfg?.memory?.qmd,
102+
limits: {
103+
...limits,
104+
maxResults: resolvePositiveIntegerConfig(limits?.maxResults, 4),
105+
maxSnippetChars: resolvePositiveIntegerConfig(limits?.maxSnippetChars, 450),
106+
maxInjectedChars: resolvePositiveIntegerConfig(limits?.maxInjectedChars, 2_200),
107+
timeoutMs: resolvePositiveIntegerConfig(limits?.timeoutMs, 4_000),
108+
},
109+
}
110+
: undefined,
111+
};
112+
},
80113
getMemorySearchManager: getMemorySearchManagerMock,
81114
readAgentMemoryFile: readAgentMemoryFileMock,
82115
}));

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

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,96 @@ describe("memory_search unavailable payloads", () => {
189189
}
190190
});
191191

192+
it("lets QMD searches run past the default tool deadline when configured", async () => {
193+
vi.useFakeTimers();
194+
try {
195+
setMemoryBackend("qmd");
196+
let searchSignal: AbortSignal | undefined;
197+
setMemorySearchImpl(async (opts) => {
198+
searchSignal = opts?.signal;
199+
return await new Promise((resolve) => {
200+
setTimeout(() => {
201+
resolve([
202+
{
203+
path: "MEMORY.md",
204+
startLine: 1,
205+
endLine: 2,
206+
score: 0.9,
207+
snippet: "slow qmd recall",
208+
source: "memory" as const,
209+
},
210+
]);
211+
}, 16_000);
212+
});
213+
});
214+
const tool = createMemorySearchToolOrThrow({
215+
config: {
216+
agents: { list: [{ id: "main", default: true }] },
217+
memory: {
218+
citations: "off",
219+
backend: "qmd",
220+
qmd: { limits: { timeoutMs: 60_000.5 } },
221+
},
222+
},
223+
});
224+
225+
const resultPromise = tool.execute("slow-qmd-search", { query: "slow recall" });
226+
await vi.advanceTimersByTimeAsync(16_000);
227+
228+
const result = await resultPromise;
229+
expect(searchSignal?.aborted).toBe(false);
230+
expect((result.details as { results?: Array<{ path: string }> }).results).toEqual([
231+
{
232+
corpus: "memory",
233+
path: "MEMORY.md",
234+
startLine: 1,
235+
endLine: 2,
236+
score: 0.9,
237+
snippet: "slow qmd recall",
238+
source: "memory",
239+
},
240+
]);
241+
} finally {
242+
vi.useRealTimers();
243+
}
244+
});
245+
246+
it("keeps a finite QMD tool deadline when the configured search does not settle", async () => {
247+
vi.useFakeTimers();
248+
try {
249+
setMemoryBackend("qmd");
250+
let settled = false;
251+
setMemorySearchImpl(async () => await new Promise(() => {}));
252+
const tool = createMemorySearchToolOrThrow({
253+
config: {
254+
agents: { list: [{ id: "main", default: true }] },
255+
memory: {
256+
backend: "qmd",
257+
qmd: { limits: { timeoutMs: 60_000 } },
258+
},
259+
},
260+
});
261+
262+
const resultPromise = tool.execute("stalled-qmd-search", { query: "slow recall" });
263+
void resultPromise.then(() => {
264+
settled = true;
265+
});
266+
267+
await vi.advanceTimersByTimeAsync(64_999);
268+
expect(settled).toBe(false);
269+
270+
await vi.advanceTimersByTimeAsync(1);
271+
const result = await resultPromise;
272+
expectUnavailableMemorySearchDetails(result.details, {
273+
error: "memory_search timed out after 65s",
274+
warning: "Memory search is unavailable due to an embedding/provider error.",
275+
action: "Check embedding provider configuration and retry memory_search.",
276+
});
277+
} finally {
278+
vi.useRealTimers();
279+
}
280+
});
281+
192282
it("keeps the timeout result when an abort-aware search rejects on abort", async () => {
193283
vi.useFakeTimers();
194284
try {

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

Lines changed: 34 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
// Memory Core plugin module implements tools behavior.
22
import { formatErrorMessage } from "openclaw/plugin-sdk/error-runtime";
3-
import type { MemorySource } from "openclaw/plugin-sdk/memory-core-host-engine-storage";
3+
import type {
4+
MemorySource,
5+
ResolvedMemoryBackendConfig,
6+
} from "openclaw/plugin-sdk/memory-core-host-engine-storage";
47
import {
58
asToolParamsRecord,
69
jsonResult,
@@ -19,6 +22,7 @@ import {
1922
resolveMemoryDreamingConfig,
2023
resolveMemoryDeepDreamingConfig,
2124
} from "openclaw/plugin-sdk/memory-core-host-status";
25+
import { resolveTimerTimeoutMs } from "openclaw/plugin-sdk/number-runtime";
2226
import { asRecord } from "./dreaming-shared.js";
2327
import { filterMemorySearchHitsBySessionVisibility } from "./session-search-visibility.js";
2428
import { recordShortTermRecalls } from "./short-term-promotion.js";
@@ -46,6 +50,7 @@ type MemoryManagerContext = Awaited<ReturnType<typeof getMemoryManagerContextWit
4650
type ActiveMemoryManagerContext = Extract<MemoryManagerContext, { manager: unknown }>;
4751

4852
const MEMORY_SEARCH_TOOL_TIMEOUT_MS = 15_000;
53+
const MEMORY_SEARCH_TOOL_QMD_TIMEOUT_OVERHEAD_MS = 5_000;
4954
const MEMORY_SEARCH_TOOL_COOLDOWN_MS = 60_000;
5055

5156
const memorySearchToolCooldowns = new Map<string, { until: number; error: string }>();
@@ -76,6 +81,23 @@ function recordMemorySearchToolCooldown(key: string, error: string): void {
7681
});
7782
}
7883

84+
function resolveMemorySearchToolTimeoutMs(params: {
85+
resolvedBackend?: ResolvedMemoryBackendConfig;
86+
requestedCorpus?: "memory" | "wiki" | "all" | "sessions";
87+
}): number {
88+
if (params.requestedCorpus === "wiki" || params.resolvedBackend?.backend !== "qmd") {
89+
return MEMORY_SEARCH_TOOL_TIMEOUT_MS;
90+
}
91+
const qmdTimeoutMs = params.resolvedBackend.qmd?.limits.timeoutMs;
92+
if (qmdTimeoutMs === undefined || qmdTimeoutMs <= MEMORY_SEARCH_TOOL_TIMEOUT_MS) {
93+
return MEMORY_SEARCH_TOOL_TIMEOUT_MS;
94+
}
95+
return resolveTimerTimeoutMs(
96+
qmdTimeoutMs + MEMORY_SEARCH_TOOL_QMD_TIMEOUT_OVERHEAD_MS,
97+
MEMORY_SEARCH_TOOL_TIMEOUT_MS,
98+
);
99+
}
100+
79101
export const testing = {
80102
resetMemorySearchToolCooldowns() {
81103
memorySearchToolCooldowns.clear();
@@ -104,8 +126,9 @@ async function runMemorySearchToolWithDeadline<T>(params: {
104126
timeoutMs: number;
105127
run: (signal: AbortSignal) => Promise<T>;
106128
}): Promise<{ status: "ok"; value: T } | { status: "unavailable"; error: string }> {
129+
const timeoutMs = resolveTimerTimeoutMs(params.timeoutMs, MEMORY_SEARCH_TOOL_TIMEOUT_MS);
107130
const timeoutError = () =>
108-
new Error(`memory_search timed out after ${Math.round(params.timeoutMs / 1000)}s`);
131+
new Error(`memory_search timed out after ${Math.round(timeoutMs / 1000)}s`);
109132
// Abort the losing task when the deadline fires so in-flight embedding work
110133
// is cancelled instead of retrying orphaned for minutes after the tool
111134
// already returned "timed out" to the agent.
@@ -118,7 +141,7 @@ async function runMemorySearchToolWithDeadline<T>(params: {
118141
// timeout result with a provider-wrapped abort error.
119142
resolve("timeout");
120143
controller.abort(timeoutError());
121-
}, params.timeoutMs);
144+
}, timeoutMs);
122145
timer.unref?.();
123146
});
124147
const task = params.run(controller.signal);
@@ -393,6 +416,9 @@ export function createMemorySearchTool(options: {
393416
});
394417
const cooldown =
395418
requestedCorpus === "wiki" ? undefined : readMemorySearchToolCooldown(cooldownKey);
419+
const memoryRuntime =
420+
requestedCorpus === "wiki" ? undefined : await loadMemoryToolRuntime();
421+
const resolvedBackend = memoryRuntime?.resolveMemoryBackendConfig({ cfg, agentId });
396422
let activeUnavailablePhase: "memory" | "supplement" | undefined;
397423
let failedUnavailablePhase: "memory" | "supplement" | undefined;
398424
const runUnavailablePhase = async <T>(
@@ -413,9 +439,11 @@ export function createMemorySearchTool(options: {
413439
};
414440

415441
const outcome = await runMemorySearchToolWithDeadline({
416-
timeoutMs: MEMORY_SEARCH_TOOL_TIMEOUT_MS,
442+
timeoutMs: resolveMemorySearchToolTimeoutMs({
443+
resolvedBackend,
444+
requestedCorpus,
445+
}),
417446
run: async (deadlineSignal) => {
418-
const { resolveMemoryBackendConfig } = await loadMemoryToolRuntime();
419447
const shouldQuerySupplements = requestedCorpus === "wiki" || requestedCorpus === "all";
420448
const shouldQueryMemory = requestedCorpus !== "wiki" && !cooldown;
421449
if (cooldown && !shouldQuerySupplements) {
@@ -555,12 +583,11 @@ export function createMemorySearchTool(options: {
555583
}
556584
const status = activeMemory.manager.status();
557585
const decorated = decorateCitations(rawResults, includeCitations);
558-
const resolved = resolveMemoryBackendConfig({ cfg, agentId });
559586
const memoryResults =
560587
status.backend === "qmd"
561588
? clampResultsByInjectedChars(
562589
decorated,
563-
resolved.qmd?.limits.maxInjectedChars,
590+
resolvedBackend?.qmd?.limits.maxInjectedChars,
564591
)
565592
: decorated;
566593
surfacedMemoryResults = memoryResults.map((result) => ({

0 commit comments

Comments
 (0)