diff --git a/extensions/memory-core/src/memory-tool-manager-mock.ts b/extensions/memory-core/src/memory-tool-manager-mock.ts index 0b906d1bf54..73a7265423c 100644 --- a/extensions/memory-core/src/memory-tool-manager-mock.ts +++ b/extensions/memory-core/src/memory-tool-manager-mock.ts @@ -8,6 +8,7 @@ type SearchImpl = (opts?: { sessionKey?: string; qmdSearchModeOverride?: "query" | "search" | "vsearch"; onDebug?: (debug: MemorySearchRuntimeDebug) => void; + signal?: AbortSignal; }) => Promise; export type MemoryReadParams = { relPath: string; from?: number; lines?: number }; type MemoryReadResult = { diff --git a/extensions/memory-core/src/memory/manager-embedding-ops.ts b/extensions/memory-core/src/memory/manager-embedding-ops.ts index 01e84fbbb41..2d489f49fc5 100644 --- a/extensions/memory-core/src/memory/manager-embedding-ops.ts +++ b/extensions/memory-core/src/memory/manager-embedding-ops.ts @@ -163,11 +163,16 @@ export function resolveMemoryIndexConcurrency(params: { export async function runEmbeddingOperationWithTimeout(params: { timeoutMs: number; message: string; + /** Caller-owned cancellation, merged with the per-call watchdog abort. */ + signal?: AbortSignal; run: (signal: AbortSignal) => Promise; }): Promise { 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(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 { + protected async embedQueryWithRetry(text: string, signal?: AbortSignal): Promise { 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"); diff --git a/extensions/memory-core/src/memory/manager-embedding-policy.test.ts b/extensions/memory-core/src/memory/manager-embedding-policy.test.ts index a2eec83b33c..93f5d14bc65 100644 --- a/extensions/memory-core/src/memory/manager-embedding-policy.test.ts +++ b/extensions/memory-core/src/memory/manager-embedding-policy.test.ts @@ -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", diff --git a/extensions/memory-core/src/memory/manager-embedding-policy.ts b/extensions/memory-core/src/memory/manager-embedding-policy.ts index 25c0326232b..a224b4abc26 100644 --- a/extensions/memory-core/src/memory/manager-embedding-policy.ts +++ b/extensions/memory-core/src/memory/manager-embedding-policy.ts @@ -125,6 +125,8 @@ export async function runMemoryEmbeddingRetryLoop(params: { waitForRetry: (delayMs: number) => Promise; maxAttempts: number; baseDelayMs: number; + /** Caller-owned cancellation; an aborted caller stops the retry loop. */ + signal?: AbortSignal; }): Promise { 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(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; diff --git a/extensions/memory-core/src/memory/manager-embedding-timeout.test.ts b/extensions/memory-core/src/memory/manager-embedding-timeout.test.ts index 61a01311432..2fe29ef31de 100644 --- a/extensions/memory-core/src/memory/manager-embedding-timeout.test.ts +++ b/extensions/memory-core/src/memory/manager-embedding-timeout.test.ts @@ -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((_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({ diff --git a/extensions/memory-core/src/memory/manager.ts b/extensions/memory-core/src/memory/manager.ts index 326e160ec89..e51dd310a0f 100644 --- a/extensions/memory-core/src/memory/manager.ts +++ b/extensions/memory-core/src/memory/manager.ts @@ -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 { 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); diff --git a/extensions/memory-core/src/memory/search-manager.test.ts b/extensions/memory-core/src/memory/search-manager.test.ts index e114fc60c99..67f25d8fcd5 100644 --- a/extensions/memory-core/src/memory/search-manager.test.ts +++ b/extensions/memory-core/src/memory/search-manager.test.ts @@ -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); diff --git a/extensions/memory-core/src/memory/search-manager.ts b/extensions/memory-core/src/memory/search-manager.ts index 7054e14fbb2..07a43466c46 100644 --- a/extensions/memory-core/src/memory/search-manager.ts +++ b/extensions/memory-core/src/memory/search-manager.ts @@ -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); diff --git a/extensions/memory-core/src/tools.test.ts b/extensions/memory-core/src/tools.test.ts index 453400d6d27..79d071f1107 100644 --- a/extensions/memory-core/src/tools.test.ts +++ b/extensions/memory-core/src/tools.test.ts @@ -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 () => { diff --git a/extensions/memory-core/src/tools.ts b/extensions/memory-core/src/tools.ts index 8c6b8252a93..946804935c7 100644 --- a/extensions/memory-core/src/tools.ts +++ b/extensions/memory-core/src/tools.ts @@ -83,14 +83,26 @@ export const testing = { async function runMemorySearchToolWithDeadline(params: { timeoutMs: number; - run: () => Promise; + run: (signal: AbortSignal) => Promise; }): 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 | 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(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); }, diff --git a/packages/memory-host-sdk/src/host/types.ts b/packages/memory-host-sdk/src/host/types.ts index 50263d5f0ab..de8edde9b5d 100644 --- a/packages/memory-host-sdk/src/host/types.ts +++ b/packages/memory-host-sdk/src/host/types.ts @@ -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; readFile(params: { relPath: string; from?: number; lines?: number }): Promise;