fix(memory-lancedb): share the recall cooldown with the auto-recall prompt-build hook (#112927)

* fix(memory-lancedb): share recall cooldown with auto-recall prompt-build hook

* fix(memory-lancedb): share recall cooldown with auto-recall

* chore: keep release changelog maintainer-owned

---------

Co-authored-by: Vincent Koc <vincentkoc@ieee.org>
This commit is contained in:
amittell
2026-07-29 13:30:41 -04:00
committed by GitHub
parent 877fa6178a
commit 9155d90002
3 changed files with 218 additions and 16 deletions

View File

@@ -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;

View File

@@ -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]);

View File

@@ -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<ReturnType<typeof runWithTimeout<MemorySearchResult[]>>>;
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;