diff --git a/extensions/memory-lancedb/embeddings.ts b/extensions/memory-lancedb/embeddings.ts index 3567f9f7ffc7..de0513139a01 100644 --- a/extensions/memory-lancedb/embeddings.ts +++ b/extensions/memory-lancedb/embeddings.ts @@ -387,6 +387,33 @@ export function formatMemoryRecallError(error: unknown): string { return error instanceof Error ? error.message : String(error); } +export function isMemoryRecallTimeoutError(error: unknown): boolean { + let current: unknown = error; + for (let depth = 0; depth < 3 && current !== undefined; depth += 1) { + const record = asRecord(current); + const name = + current instanceof Error ? current.name : typeof record?.name === "string" ? record.name : ""; + const message = + current instanceof Error + ? current.message + : typeof record?.message === "string" + ? record.message + : ""; + const code = typeof record?.code === "string" ? record.code : ""; + if ( + name === "APIConnectionTimeoutError" || + name === "TimeoutError" || + code === "ETIMEDOUT" || + /^UND_ERR_.*_TIMEOUT$/.test(code) || + /\btimed out\b/i.test(message) + ) { + return true; + } + current = record?.cause; + } + return false; +} + export function buildMemoryRecallUnavailableResult(error: string): AgentToolResult<{ count: number; disabled: true; @@ -413,6 +440,7 @@ export class MemoryRecallEmbeddingError extends Error { export const testing = { isEmbeddingDimensionsRejectedError, + isMemoryRecallTimeoutError, runWithTimeout, truncateEmbeddingVector, } as const; diff --git a/extensions/memory-lancedb/index.test.ts b/extensions/memory-lancedb/index.test.ts index 344d9335da76..72a00d2a2763 100644 --- a/extensions/memory-lancedb/index.test.ts +++ b/extensions/memory-lancedb/index.test.ts @@ -1305,7 +1305,7 @@ describe("memory plugin e2e", () => { }); }); - test("bounds auto-recall latency during prompt build", async () => { + test("shares only embedding timeout cooldown across recall paths", async () => { vi.useFakeTimers(); const post = vi.fn( () => @@ -1320,14 +1320,14 @@ describe("memory plugin e2e", () => { }), ); const ensureGlobalUndiciEnvProxyDispatcher = vi.fn(); + const toArray = vi.fn(() => new Promise(() => {})); + const limit = vi.fn(() => ({ toArray })); const loadLanceDbModule = vi.fn(async () => ({ connect: vi.fn(async () => ({ tableNames: vi.fn(async () => ["memories"]), openTable: vi.fn(async () => ({ schema: createAgentScopedSchemaMock(), - vectorSearch: vi.fn(() => - createAgentScopedVectorQuery(vi.fn(() => ({ toArray: vi.fn(async () => []) }))), - ), + vectorSearch: vi.fn(() => createAgentScopedVectorQuery(limit)), countRows: vi.fn(async () => 0), add: vi.fn(async () => undefined), delete: vi.fn(async () => undefined), @@ -1342,6 +1342,7 @@ describe("memory plugin e2e", () => { loadLanceDbModule, run: async (dynamicMemoryPlugin) => { const on = vi.fn(); + const registeredTools: any[] = []; const logger = { info: vi.fn(), warn: vi.fn(), @@ -1364,7 +1365,9 @@ describe("memory plugin e2e", () => { }, runtime: {}, logger, - registerTool: vi.fn(), + registerTool: (tool: any, opts: any) => { + registeredTools.push({ tool, opts }); + }, registerCli: vi.fn(), registerService: vi.fn(), on, @@ -1378,10 +1381,8 @@ describe("memory plugin e2e", () => { )?.[1]; expect(beforePromptBuild).toBeTypeOf("function"); - const resultPromise = beforePromptBuild?.( - { prompt: "what editor should i use?", messages: [] }, - { agentId: "main" }, - ); + const hookEvent = { prompt: "what editor should i use?", messages: [] }; + const resultPromise = beforePromptBuild?.(hookEvent, { agentId: "main" }); await vi.advanceTimersByTimeAsync(15_000); await expect(resultPromise).resolves.toBeUndefined(); @@ -1394,6 +1395,110 @@ describe("memory plugin e2e", () => { expect(logger.warn).toHaveBeenCalledWith( "memory-lancedb: auto-recall timed out after 15000ms; skipping memory injection to avoid stalling agent startup", ); + + expect(await beforePromptBuild?.(hookEvent, { agentId: "main" })).toBeUndefined(); + expect(post).toHaveBeenCalledTimes(1); + expect(logger.debug).toHaveBeenCalledWith( + "memory-lancedb: auto-recall skipped during recall cooldown: auto-recall timed out after 15s", + ); + + const recallTool = materializeRegisteredTool( + registeredTools.find((tool) => tool.opts?.name === "memory_recall")?.tool, + ); + if (!recallTool) { + throw new Error("memory_recall tool was not registered"); + } + const toolResult = await recallTool.execute("cooldown-call", { query: "editor" }); + expect(toolResult.details).toMatchObject({ + count: 0, + disabled: true, + unavailable: true, + error: "auto-recall timed out after 15s", + }); + expect(post).toHaveBeenCalledTimes(1); + + await vi.advanceTimersByTimeAsync(60_000); + const sdkTimeoutError = Object.assign(new Error("Request timed out."), { + name: "APIConnectionTimeoutError", + }); + post.mockRejectedValueOnce(sdkTimeoutError); + await expect( + beforePromptBuild?.(hookEvent, { agentId: "main" }), + ).resolves.toBeUndefined(); + expect(post).toHaveBeenCalledTimes(2); + + const sdkTimeoutToolResult = await recallTool.execute("sdk-timeout-cooldown-call", { + query: "editor", + }); + expect(sdkTimeoutToolResult.details).toMatchObject({ + count: 0, + disabled: true, + unavailable: true, + error: "Request timed out.", + }); + expect(post).toHaveBeenCalledTimes(2); + + await vi.advanceTimersByTimeAsync(60_000); + post.mockResolvedValueOnce({ data: [{ embedding: [0.1, 0.2, 0.3] }] }); + const probeResult = beforePromptBuild?.(hookEvent, { agentId: "main" }); + await vi.advanceTimersByTimeAsync(0); + expect(loadLanceDbModule).toHaveBeenCalledTimes(1); + await vi.advanceTimersByTimeAsync(15_000); + await expect(probeResult).resolves.toBeUndefined(); + expect(post).toHaveBeenCalledTimes(3); + + post.mockRejectedValueOnce(Object.assign(new Error("bad auto query"), { status: 400 })); + const retryResult = beforePromptBuild?.(hookEvent, { agentId: "main" }); + await vi.advanceTimersByTimeAsync(0); + expect(post).toHaveBeenCalledTimes(4); + await expect(retryResult).resolves.toBeUndefined(); + + post.mockRejectedValueOnce(Object.assign(new Error("bad tool query"), { status: 400 })); + const toolErrorResult = await recallTool.execute("error-call", { query: "editor" }); + expect(toolErrorResult.details).toMatchObject({ + count: 0, + disabled: true, + unavailable: true, + error: "bad tool query", + }); + expect(post).toHaveBeenCalledTimes(5); + + post.mockRejectedValueOnce(sdkTimeoutError); + const toolSdkTimeoutResult = await recallTool.execute("sdk-timeout-call", { + query: "editor", + }); + expect(toolSdkTimeoutResult.details).toMatchObject({ + count: 0, + disabled: true, + unavailable: true, + error: "Request timed out.", + }); + expect(post).toHaveBeenCalledTimes(6); + + expect(await beforePromptBuild?.(hookEvent, { agentId: "main" })).toBeUndefined(); + expect(post).toHaveBeenCalledTimes(6); + + await vi.advanceTimersByTimeAsync(60_000); + + post.mockResolvedValueOnce({ data: [{ embedding: [0.1, 0.2, 0.3] }] }); + const toolSearchResult = recallTool.execute("search-timeout-call", { query: "editor" }); + await vi.advanceTimersByTimeAsync(0); + expect(post).toHaveBeenCalledTimes(7); + await vi.advanceTimersByTimeAsync(15_000); + await expect(toolSearchResult).resolves.toMatchObject({ + details: { + count: 0, + disabled: true, + unavailable: true, + error: "memory_recall timed out after 15s", + }, + }); + + const finalResult = beforePromptBuild?.(hookEvent, { agentId: "main" }); + await vi.advanceTimersByTimeAsync(0); + expect(post).toHaveBeenCalledTimes(8); + await vi.advanceTimersByTimeAsync(15_000); + await expect(finalResult).resolves.toBeUndefined(); await vi.advanceTimersByTimeAsync(15_000); }, }); @@ -3333,6 +3438,38 @@ describe("memory plugin e2e", () => { ).toBe(false); }); + test("recognizes embedding timeout errors without classifying fast request failures", () => { + expect( + testing.isMemoryRecallTimeoutError( + Object.assign(new Error("Request timed out."), { + name: "APIConnectionTimeoutError", + }), + ), + ).toBe(true); + expect( + testing.isMemoryRecallTimeoutError( + Object.assign(new Error("socket deadline"), { code: "ETIMEDOUT" }), + ), + ).toBe(true); + expect( + testing.isMemoryRecallTimeoutError( + Object.assign(new Error("provider aborted"), { + cause: new Error("memory-lancedb embedding timed out"), + }), + ), + ).toBe(true); + expect( + testing.isMemoryRecallTimeoutError( + Object.assign(new Error("headers deadline"), { code: "UND_ERR_HEADERS_TIMEOUT" }), + ), + ).toBe(true); + expect( + testing.isMemoryRecallTimeoutError( + Object.assign(new Error("bad request"), { status: 400, code: "invalid_request_error" }), + ), + ).toBe(false); + }); + test("normalizes locally truncated embeddings and rejects short vectors", () => { expect(testing.truncateEmbeddingVector([3, 4, 12], 2, "test-model")).toEqual([0.6, 0.8]); expect(testing.truncateEmbeddingVector([0, 0, 1], 2, "test-model")).toEqual([0, 0]); diff --git a/extensions/memory-lancedb/index.ts b/extensions/memory-lancedb/index.ts index 2a97182f3181..6538d02f3c9f 100644 --- a/extensions/memory-lancedb/index.ts +++ b/extensions/memory-lancedb/index.ts @@ -25,6 +25,7 @@ import { buildMemoryRecallUnavailableResult, createEmbeddings, formatMemoryRecallError, + isMemoryRecallTimeoutError, MemoryRecallEmbeddingError, runWithTimeout, } from "./embeddings.js"; @@ -53,7 +54,7 @@ const loadMemoryHostCoreModule = createLazyRuntimeModule( const DEFAULT_AUTO_RECALL_TIMEOUT_MS = 15_000; const DEFAULT_TOOL_RECALL_TIMEOUT_MS = 15_000; -const DEFAULT_TOOL_RECALL_COOLDOWN_MS = 60_000; +const DEFAULT_RECALL_COOLDOWN_MS = 60_000; const DEFAULT_TOOL_RECALL_OVERFETCH_EXTRA = 10; // Auto-recall over-fetches from the vector store, then filters envelope sludge @@ -178,7 +179,7 @@ export default definePluginEntry({ }; const recordMemoryRecallCooldown = (agentId: string, error: string): void => { memoryRecallCooldowns.set(agentId, { - until: Date.now() + DEFAULT_TOOL_RECALL_COOLDOWN_MS, + until: Date.now() + DEFAULT_RECALL_COOLDOWN_MS, error, }); }; @@ -221,6 +222,7 @@ export default definePluginEntry({ if (cooldown) { return buildMemoryRecallUnavailableResult(cooldown.error); } + let recallPhase: "embedding" | "search" = "embedding"; let recall: Awaited>>; try { recall = await runWithTimeout({ @@ -235,6 +237,7 @@ export default definePluginEntry({ } catch (error) { throw new MemoryRecallEmbeddingError(error); } + recallPhase = "search"; return await db.search( agentId, vector, @@ -248,7 +251,9 @@ export default definePluginEntry({ throw error; } const message = formatMemoryRecallError(error.originalError); - recordMemoryRecallCooldown(agentId, message); + if (isMemoryRecallTimeoutError(error.originalError)) { + recordMemoryRecallCooldown(agentId, message); + } api.logger.warn?.( `memory-lancedb: memory_recall failed: ${message}; returning unavailable memory result`, ); @@ -256,7 +261,9 @@ export default definePluginEntry({ } if (recall.status === "timeout") { const message = `memory_recall timed out after ${Math.round(DEFAULT_TOOL_RECALL_TIMEOUT_MS / 1000)}s`; - recordMemoryRecallCooldown(agentId, message); + if (recallPhase === "embedding") { + recordMemoryRecallCooldown(agentId, message); + } api.logger.warn?.( `memory-lancedb: memory_recall timed out after ${DEFAULT_TOOL_RECALL_TIMEOUT_MS}ms; returning unavailable memory result`, ); @@ -502,6 +509,15 @@ export default definePluginEntry({ if (!event.prompt || event.prompt.length < 5) { return undefined; } + // One hung embedding request must not stall both automatic and explicit recall. + // Keep the breaker per agent so unrelated memory namespaces still probe. + const cooldown = readMemoryRecallCooldown(agentId); + if (cooldown) { + api.logger.debug?.( + `memory-lancedb: auto-recall skipped during recall cooldown: ${cooldown.error}`, + ); + return undefined; + } try { const recallQuery = normalizeRecallQuery( @@ -514,18 +530,33 @@ export default definePluginEntry({ if (!recallQuery) { return undefined; } + let recallPhase: "embedding" | "search" = "embedding"; const recall = await runWithTimeout({ timeoutMs: DEFAULT_AUTO_RECALL_TIMEOUT_MS, task: async () => { - const vector = await embeddings.embed(recallQuery, { - timeoutMs: DEFAULT_AUTO_RECALL_TIMEOUT_MS, - }); + let vector: number[]; + try { + vector = await embeddings.embed(recallQuery, { + timeoutMs: DEFAULT_AUTO_RECALL_TIMEOUT_MS, + }); + } catch (error) { + throw new MemoryRecallEmbeddingError(error); + } + // Keep one end-to-end deadline, but only let embedding timeouts trip + // the shared breaker. LanceDB stalls remain retryable next turn. + recallPhase = "search"; // Overfetch to compensate for sludge filtering: if contaminated // entries occupy the top slots we still surface enough clean ones. return await db.search(agentId, vector, DEFAULT_AUTO_RECALL_OVERFETCH_LIMIT, 0.3); }, }); if (recall.status === "timeout") { + if (recallPhase === "embedding") { + recordMemoryRecallCooldown( + agentId, + `auto-recall timed out after ${Math.round(DEFAULT_AUTO_RECALL_TIMEOUT_MS / 1000)}s`, + ); + } api.logger.warn?.( `memory-lancedb: auto-recall timed out after ${DEFAULT_AUTO_RECALL_TIMEOUT_MS}ms; skipping memory injection to avoid stalling agent startup`, ); @@ -552,6 +583,12 @@ export default definePluginEntry({ prependContext: context, }; } catch (err) { + if ( + err instanceof MemoryRecallEmbeddingError && + isMemoryRecallTimeoutError(err.originalError) + ) { + recordMemoryRecallCooldown(agentId, formatMemoryRecallError(err.originalError)); + } api.logger.warn(`memory-lancedb: recall failed: ${String(err)}`); } return undefined;