fix(memory): abort orphaned embedding work when memory_search times out (#91742)

* fix(memory): abort orphaned embedding work when memory_search times out

memory_search raced its 15s deadline with Promise.race and returned a clean
timeout to the agent, but the underlying embedQueryWithRetry loop kept
retrying (3 attempts x 60s) against the embedding backend with no consumer.
Thread the tool-owned AbortSignal through manager.search ->
embedQueryWithRetry -> runEmbeddingOperationWithTimeout so the deadline
cancels in-flight embedding work, stops the retry loop, and skips
fallback-provider activation for an absent caller.

Fixes #91718

* fix(memory): let the deadline result win before aborting the search

Abort listeners dispatch synchronously, so an abort-aware search could
reject the raced task before the timeout promise resolved and replace the
stable 'memory_search timed out after 15s' result with a provider-wrapped
abort error. Resolve the timeout first, then abort.

* fix(memory): scope deadline abort to builtin embeddings

* fix(memory): preserve deadline signal across fallback

---------

Co-authored-by: Vincent Koc <vincentkoc@ieee.org>
This commit is contained in:
Dream Hunter 2026-06-11 19:36:11 +08:00 committed by GitHub
parent 4f3c2cd2df
commit 8d72cb9401
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 137 additions and 13 deletions

View file

@ -8,6 +8,7 @@ type SearchImpl = (opts?: {
sessionKey?: string;
qmdSearchModeOverride?: "query" | "search" | "vsearch";
onDebug?: (debug: MemorySearchRuntimeDebug) => void;
signal?: AbortSignal;
}) => Promise<unknown[]>;
export type MemoryReadParams = { relPath: string; from?: number; lines?: number };
type MemoryReadResult = {

View file

@ -163,11 +163,16 @@ export function resolveMemoryIndexConcurrency(params: {
export async function runEmbeddingOperationWithTimeout<T>(params: {
timeoutMs: number;
message: string;
/** Caller-owned cancellation, merged with the per-call watchdog abort. */
signal?: AbortSignal;
run: (signal: AbortSignal) => Promise<T>;
}): Promise<T> {
const controller = new AbortController();
const signal = params.signal
? AbortSignal.any([params.signal, controller.signal])
: controller.signal;
if (!Number.isFinite(params.timeoutMs) || params.timeoutMs <= 0) {
return await params.run(controller.signal);
return await params.run(signal);
}
const timeoutMs = resolveTimerTimeoutMs(params.timeoutMs, 1);
let timer: NodeJS.Timeout | null = null;
@ -179,7 +184,7 @@ export async function runEmbeddingOperationWithTimeout<T>(params: {
}, timeoutMs);
});
try {
const operation = params.run(controller.signal);
const operation = params.run(signal);
return (await Promise.race([operation, timeoutPromise])) as T;
} finally {
if (timer) {
@ -493,7 +498,7 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps {
});
}
protected async embedQueryWithRetry(text: string): Promise<number[]> {
protected async embedQueryWithRetry(text: string, signal?: AbortSignal): Promise<number[]> {
const provider = this.provider;
if (!provider) {
throw new Error("Cannot embed query in FTS-only mode (no embedding provider)");
@ -501,14 +506,17 @@ export abstract class MemoryManagerEmbeddingOps extends MemoryManagerSyncOps {
try {
return await runMemoryEmbeddingRetryLoop({
run: async () => {
signal?.throwIfAborted();
const timeoutMs = this.resolveEmbeddingTimeout("query");
log.debug("memory embeddings: query start", { provider: provider.id, timeoutMs });
return await runEmbeddingOperationWithTimeout({
timeoutMs,
message: `memory embeddings query timed out after ${Math.round(timeoutMs / 1000)}s`,
run: async (signal) => await provider.embedQuery(text, { signal }),
signal,
run: async (opSignal) => await provider.embedQuery(text, { signal: opSignal }),
});
},
signal,
isRetryable: isRetryableMemoryEmbeddingError,
waitForRetry: async (delayMs) => {
await this.waitForEmbeddingRetry(delayMs, "retrying query");

View file

@ -75,6 +75,30 @@ describe("memory embedding policy", () => {
expect(waits).toEqual([500, 1000]);
});
it("stops retrying after the caller signal aborts, even for retryable-looking errors", async () => {
const controller = new AbortController();
const run = vi.fn(async () => {
controller.abort(new Error("memory_search timed out after 15s"));
// "timed out" matches the retryable transport pattern; abort must still win.
throw new Error("memory embeddings query timed out after 60s");
});
const waitForRetry = vi.fn(async () => {});
await expect(
runMemoryEmbeddingRetryLoop({
run,
isRetryable: isRetryableMemoryEmbeddingError,
waitForRetry,
maxAttempts: 3,
baseDelayMs: 500,
signal: controller.signal,
}),
).rejects.toThrow("memory embeddings query timed out after 60s");
expect(run).toHaveBeenCalledTimes(1);
expect(waitForRetry).not.toHaveBeenCalled();
});
it("retries transient socket/network embedding errors", () => {
const splittableMessages = [
"TypeError: fetch failed | other side closed",

View file

@ -125,6 +125,8 @@ export async function runMemoryEmbeddingRetryLoop<T>(params: {
waitForRetry: (delayMs: number) => Promise<void>;
maxAttempts: number;
baseDelayMs: number;
/** Caller-owned cancellation; an aborted caller stops the retry loop. */
signal?: AbortSignal;
}): Promise<T> {
const attempts = Math.max(1, params.maxAttempts);
for (const attempt of Array.from({ length: attempts }, (_, index) => index + 1)) {
@ -132,6 +134,12 @@ export async function runMemoryEmbeddingRetryLoop<T>(params: {
try {
return await params.run();
} catch (err) {
// Abort must win over retryable-looking failures: abort reasons often
// carry "timed out" messages that match the retryable transport
// patterns and would otherwise keep retrying for an absent caller.
if (params.signal?.aborted) {
throw err;
}
const message = formatErrorMessage(err);
if (!params.isRetryable(message) || attempt >= params.maxAttempts) {
throw err;

View file

@ -135,6 +135,31 @@ describe("memory embedding timeout abort", () => {
expect(signalSeen?.aborted).toBe(true);
});
it("aborts the provider operation when the caller signal aborts before the watchdog", async () => {
const external = new AbortController();
let signalSeen: AbortSignal | undefined;
const resultPromise = runEmbeddingOperationWithTimeout({
timeoutMs: 60_000,
message: "memory embeddings query timed out after 60s",
signal: external.signal,
run: async (signal) => {
signalSeen = signal;
return await new Promise<number[]>((_resolve, reject) => {
signal.addEventListener(
"abort",
() => reject(toLintErrorObject(signal.reason, "Non-Error rejection")),
{ once: true },
);
});
},
});
external.abort(new Error("memory_search timed out after 15s"));
await expect(resultPromise).rejects.toThrow("memory_search timed out after 15s");
expect(signalSeen?.aborted).toBe(true);
});
it("keeps the timeout error when a provider abort listener rejects generically", async () => {
vi.useFakeTimers();
const resultPromise = runEmbeddingOperationWithTimeout({

View file

@ -592,6 +592,8 @@ export class MemoryIndexManager extends MemoryManagerEmbeddingOps implements Mem
onDebug?: (debug: MemorySearchRuntimeDebug) => void;
/** When set, only these chunk sources are considered (must be enabled for this manager). */
sources?: MemorySource[];
/** Caller-owned cancellation; aborts in-flight embedding work when the caller stops waiting. */
signal?: AbortSignal;
},
): Promise<MemorySearchResult[]> {
opts?.onDebug?.({ backend: "builtin" });
@ -758,8 +760,13 @@ export class MemoryIndexManager extends MemoryManagerEmbeddingOps implements Mem
let queryVec: number[];
try {
queryVec = await this.embedQueryWithRetry(cleaned);
queryVec = await this.embedQueryWithRetry(cleaned, opts?.signal);
} catch (err) {
// An aborted caller already stopped waiting; skip fallback-provider
// activation so the abandoned search stops instead of re-embedding.
if (opts?.signal?.aborted) {
throw err;
}
const message = formatErrorMessage(err);
const activatedFallback = this.shouldFallbackOnError(err)
? await this.activateFallbackProvider(message).catch((fallbackErr: unknown) => {
@ -778,7 +785,7 @@ export class MemoryIndexManager extends MemoryManagerEmbeddingOps implements Mem
return [];
}
keywordResults = await loadKeywordResults();
queryVec = await this.embedQueryWithRetry(cleaned);
queryVec = await this.embedQueryWithRetry(cleaned, opts?.signal);
} else if (!this.provider && this.fts.enabled && this.fts.available) {
log.warn(`memory search: embeddings unavailable; using keyword-only results: ${message}`);
return this.selectScoredResults(keywordResults, maxResults, minScore, 0);

View file

@ -339,9 +339,11 @@ describe("getMemorySearchManager caching", () => {
errorMessage: "qmd query failed",
});
const fallbackResults = await firstManager.search("hello");
const controller = new AbortController();
const fallbackResults = await firstManager.search("hello", { signal: controller.signal });
expect(fallbackResults).toHaveLength(1);
expect(fallbackResults[0]?.path).toBe("MEMORY.md");
expect(fallbackSearch).toHaveBeenCalledWith("hello", { signal: controller.signal });
const second = await getMemorySearchManager({ cfg, agentId: retryAgentId });
requireManager(second);

View file

@ -348,6 +348,7 @@ class BorrowedMemoryManager implements MemorySearchManager {
qmdSearchModeOverride?: "query" | "search" | "vsearch";
onDebug?: (debug: MemorySearchRuntimeDebug) => void;
sources?: MemorySource[];
signal?: AbortSignal;
},
) {
return await this.inner.search(query, opts);

View file

@ -157,8 +157,10 @@ describe("memory_search unavailable payloads", () => {
vi.useFakeTimers();
try {
let searchCalls = 0;
setMemorySearchImpl(async () => {
let searchSignal: AbortSignal | undefined;
setMemorySearchImpl(async (opts) => {
searchCalls += 1;
searchSignal = opts?.signal;
return await new Promise(() => {});
});
const tool = createMemorySearchToolOrThrow();
@ -172,6 +174,8 @@ describe("memory_search unavailable payloads", () => {
warning: "Memory search is unavailable due to an embedding/provider error.",
action: "Check embedding provider configuration and retry memory_search.",
});
// The deadline must abort the orphaned search, not just race past it.
expect(searchSignal?.aborted).toBe(true);
const cooldownResult = await tool.execute("search-cooldown", { query: "hello again" });
expectUnavailableMemorySearchDetails(cooldownResult.details, {
error: "memory_search timed out after 15s",
@ -184,6 +188,35 @@ describe("memory_search unavailable payloads", () => {
}
});
it("keeps the timeout result when an abort-aware search rejects on abort", async () => {
vi.useFakeTimers();
try {
setMemorySearchImpl(
async (opts) =>
await new Promise((_resolve, reject) => {
opts?.signal?.addEventListener(
"abort",
() => reject(new Error("openai-compatible embeddings query failed: aborted")),
{ once: true },
);
}),
);
const tool = createMemorySearchToolOrThrow();
const resultPromise = tool.execute("abort-aware-timeout", { query: "hello" });
await vi.advanceTimersByTimeAsync(15_000);
const result = await resultPromise;
expectUnavailableMemorySearchDetails(result.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.",
});
} finally {
vi.useRealTimers();
}
});
it("re-resolves the manager once when a cached sqlite handle was closed", async () => {
let searchCalls = 0;
setMemorySearchImpl(async () => {

View file

@ -83,14 +83,26 @@ export const testing = {
async function runMemorySearchToolWithDeadline<T>(params: {
timeoutMs: number;
run: () => Promise<T>;
run: (signal: AbortSignal) => Promise<T>;
}): Promise<{ status: "ok"; value: T } | { status: "unavailable"; error: string }> {
const timeoutError = () =>
new Error(`memory_search timed out after ${Math.round(params.timeoutMs / 1000)}s`);
// Abort the losing task when the deadline fires so in-flight embedding work
// is cancelled instead of retrying orphaned for minutes after the tool
// already returned "timed out" to the agent.
const controller = new AbortController();
let timer: ReturnType<typeof setTimeout> | undefined;
const timeoutPromise = new Promise<"timeout">((resolve) => {
timer = setTimeout(() => resolve("timeout"), params.timeoutMs);
timer = setTimeout(() => {
// Resolve before aborting: abort listeners run synchronously and an
// abort-aware search could reject the task first, replacing the stable
// timeout result with a provider-wrapped abort error.
resolve("timeout");
controller.abort(timeoutError());
}, params.timeoutMs);
timer.unref?.();
});
const task = params.run();
const task = params.run(controller.signal);
task.catch(() => undefined);
try {
@ -98,7 +110,7 @@ async function runMemorySearchToolWithDeadline<T>(params: {
if (result === "timeout") {
return {
status: "unavailable",
error: `memory_search timed out after ${Math.round(params.timeoutMs / 1000)}s`,
error: timeoutError().message,
};
}
return { status: "ok", value: result as T };
@ -382,7 +394,7 @@ export function createMemorySearchTool(options: {
const outcome = await runMemorySearchToolWithDeadline({
timeoutMs: MEMORY_SEARCH_TOOL_TIMEOUT_MS,
run: async () => {
run: async (deadlineSignal) => {
const { resolveMemoryBackendConfig } = await loadMemoryToolRuntime();
const shouldQuerySupplements = requestedCorpus === "wiki" || requestedCorpus === "all";
const shouldQueryMemory = requestedCorpus !== "wiki" && !cooldown;
@ -454,6 +466,7 @@ export function createMemorySearchTool(options: {
minScore,
sessionKey: options.agentSessionKey,
qmdSearchModeOverride,
signal: deadlineSignal,
onDebug: (debug: MemorySearchRuntimeDebug) => {
runtimeDebug.push(debug);
},

View file

@ -101,6 +101,8 @@ export interface MemorySearchManager {
qmdSearchModeOverride?: "query" | "search" | "vsearch";
onDebug?: (debug: MemorySearchRuntimeDebug) => void;
sources?: MemorySource[];
/** Optional caller cancellation; managers consume it where their runtime supports cancellation. */
signal?: AbortSignal;
},
): Promise<MemorySearchResult[]>;
readFile(params: { relPath: string; from?: number; lines?: number }): Promise<MemoryReadResult>;