mirror of
https://github.com/openclaw/openclaw.git
synced 2026-08-20 06:15:55 +00:00
refactor(memory): simplify LanceDB embedding lifecycle (#125841)
This commit is contained in:
Peter Steinberger
GitHub
parent
5169760bde
commit
2c5a694ccb
3 changed files with 172 additions and 313 deletions
|
|
@ -4,6 +4,8 @@ import {
|
|||
ensureAuthProfileStore,
|
||||
replaceRuntimeAuthProfileStoreSnapshots,
|
||||
} from "openclaw/plugin-sdk/agent-runtime";
|
||||
import { createDeferred } from "openclaw/plugin-sdk/extension-shared";
|
||||
import type { MemoryEmbeddingProvider } from "openclaw/plugin-sdk/memory-core-host-engine-embeddings";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import type { OpenClawPluginApi } from "./api.js";
|
||||
import type { MemoryConfig } from "./config.js";
|
||||
|
|
@ -61,6 +63,27 @@ function embed(
|
|||
return embeddings.embed(agentId, text, embedding);
|
||||
}
|
||||
|
||||
function providerResult(
|
||||
params: {
|
||||
id?: string;
|
||||
model?: string;
|
||||
vector?: number[];
|
||||
embedQuery?: MemoryEmbeddingProvider["embedQuery"];
|
||||
close?: NonNullable<MemoryEmbeddingProvider["close"]>;
|
||||
} = {},
|
||||
) {
|
||||
const vector = params.vector ?? [0.1, 0.2, 0.3];
|
||||
return {
|
||||
provider: {
|
||||
id: params.id ?? "openai",
|
||||
model: params.model ?? "text-embedding-3-small",
|
||||
embedQuery: params.embedQuery ?? vi.fn(async () => vector),
|
||||
embedBatch: vi.fn(async () => [vector]),
|
||||
...(params.close ? { close: params.close } : {}),
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
describe("memory-lancedb provider lifecycle", () => {
|
||||
it("authenticates private agent embeddings without using the default agent's credentials", async () => {
|
||||
const config = {};
|
||||
|
|
@ -70,14 +93,7 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
if (options.agentDir !== "/tmp/agent-private") {
|
||||
throw new Error("No provider credential for the default agent");
|
||||
}
|
||||
return {
|
||||
provider: {
|
||||
id: "openai",
|
||||
model: "text-embedding-3-small",
|
||||
embedQuery,
|
||||
embedBatch: vi.fn(async () => [[0.1, 0.2, 0.3]]),
|
||||
},
|
||||
};
|
||||
return providerResult({ embedQuery });
|
||||
});
|
||||
providerMocks.getMemoryEmbeddingProvider.mockReturnValue({
|
||||
id: "openai",
|
||||
|
|
@ -110,20 +126,15 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
const closedAgentDirs: string[] = [];
|
||||
const createProvider = vi.fn(async (options: { agentDir?: string }) => {
|
||||
const agentDir = options.agentDir ?? "unscoped";
|
||||
return {
|
||||
provider: {
|
||||
id: "openai",
|
||||
model: "text-embedding-3-small",
|
||||
embedQuery: vi.fn(async (text: string) => {
|
||||
requests.push({ agentDir, text });
|
||||
return [0.1, 0.2, 0.3];
|
||||
}),
|
||||
embedBatch: vi.fn(async () => [[0.1, 0.2, 0.3]]),
|
||||
close: vi.fn(async () => {
|
||||
closedAgentDirs.push(agentDir);
|
||||
}),
|
||||
},
|
||||
};
|
||||
return providerResult({
|
||||
embedQuery: vi.fn(async (text: string) => {
|
||||
requests.push({ agentDir, text });
|
||||
return [0.1, 0.2, 0.3];
|
||||
}),
|
||||
close: vi.fn(async () => {
|
||||
closedAgentDirs.push(agentDir);
|
||||
}),
|
||||
});
|
||||
});
|
||||
providerMocks.getMemoryEmbeddingProvider.mockReturnValue({
|
||||
id: "openai",
|
||||
|
|
@ -164,17 +175,11 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
const closedAgentDirs: string[] = [];
|
||||
const createProvider = vi.fn(async (options: { agentDir?: string }) => {
|
||||
const agentDir = options.agentDir ?? "unscoped";
|
||||
return {
|
||||
provider: {
|
||||
id: "openai",
|
||||
model: "text-embedding-3-small",
|
||||
embedQuery: vi.fn(async () => [0.1, 0.2, 0.3]),
|
||||
embedBatch: vi.fn(async () => [[0.1, 0.2, 0.3]]),
|
||||
close: vi.fn(async () => {
|
||||
closedAgentDirs.push(agentDir);
|
||||
}),
|
||||
},
|
||||
};
|
||||
return providerResult({
|
||||
close: vi.fn(async () => {
|
||||
closedAgentDirs.push(agentDir);
|
||||
}),
|
||||
});
|
||||
});
|
||||
providerMocks.getMemoryEmbeddingProvider.mockReturnValue({
|
||||
id: "openai",
|
||||
|
|
@ -249,18 +254,13 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
throw new Error("Private agent credentials were revoked");
|
||||
}
|
||||
const credential = profile.key;
|
||||
return {
|
||||
provider: {
|
||||
id: "openai",
|
||||
model: "text-embedding-3-small",
|
||||
embedQuery: vi.fn(async (text: string) => {
|
||||
requests.push({ text, credential });
|
||||
return [0.1, 0.2, 0.3];
|
||||
}),
|
||||
embedBatch: vi.fn(async () => [[0.1, 0.2, 0.3]]),
|
||||
close: closeProvider,
|
||||
},
|
||||
};
|
||||
return providerResult({
|
||||
embedQuery: vi.fn(async (text: string) => {
|
||||
requests.push({ text, credential });
|
||||
return [0.1, 0.2, 0.3];
|
||||
}),
|
||||
close: closeProvider,
|
||||
});
|
||||
});
|
||||
providerMocks.getMemoryEmbeddingProvider.mockReturnValue({
|
||||
id: "openai",
|
||||
|
|
@ -336,18 +336,13 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
throw new Error("Inherited main credential is unavailable");
|
||||
}
|
||||
const credential = profile.key;
|
||||
return {
|
||||
provider: {
|
||||
id: "openai",
|
||||
model: "text-embedding-3-small",
|
||||
embedQuery: vi.fn(async (text: string) => {
|
||||
requests.push({ agentDir, credential, text });
|
||||
return [0.1, 0.2, 0.3];
|
||||
}),
|
||||
embedBatch: vi.fn(async () => [[0.1, 0.2, 0.3]]),
|
||||
close: async () => closeProvider(agentDir, credential),
|
||||
},
|
||||
};
|
||||
return providerResult({
|
||||
embedQuery: vi.fn(async (text: string) => {
|
||||
requests.push({ agentDir, credential, text });
|
||||
return [0.1, 0.2, 0.3];
|
||||
}),
|
||||
close: async () => closeProvider(agentDir, credential),
|
||||
});
|
||||
});
|
||||
providerMocks.getMemoryEmbeddingProvider.mockReturnValue({
|
||||
id: "openai",
|
||||
|
|
@ -434,15 +429,7 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
if (options.config === revokedConfig) {
|
||||
throw new Error("Private agent credentials were revoked");
|
||||
}
|
||||
return {
|
||||
provider: {
|
||||
id: "openai",
|
||||
model: "text-embedding-3-small",
|
||||
embedQuery: oldEmbedQuery,
|
||||
embedBatch: vi.fn(async () => [[0.1, 0.2, 0.3]]),
|
||||
close: closeOldProvider,
|
||||
},
|
||||
};
|
||||
return providerResult({ embedQuery: oldEmbedQuery, close: closeOldProvider });
|
||||
});
|
||||
providerMocks.getMemoryEmbeddingProvider.mockReturnValue({
|
||||
id: "openai",
|
||||
|
|
@ -472,34 +459,25 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
});
|
||||
|
||||
it("rotates live embedding overrides after admitted work drains", async () => {
|
||||
let releaseOldEmbedding: () => void = () => {};
|
||||
const oldEmbeddingGate = new Promise<void>((resolve) => {
|
||||
releaseOldEmbedding = resolve;
|
||||
});
|
||||
let oldEmbeddingStarted: () => void = () => {};
|
||||
const oldEmbeddingStart = new Promise<void>((resolve) => {
|
||||
oldEmbeddingStarted = resolve;
|
||||
});
|
||||
const oldEmbeddingGate = createDeferred<void>();
|
||||
const oldEmbeddingStart = createDeferred<void>();
|
||||
const closeOldProvider = vi.fn(async () => {});
|
||||
const closeReplacementProvider = vi.fn(async () => {});
|
||||
const createProvider = vi.fn(
|
||||
async (options: { provider: string; model: string; remote?: { apiKey?: string } }) => {
|
||||
const isOld = options.remote?.apiKey === "fixture-old-key";
|
||||
return {
|
||||
provider: {
|
||||
id: options.provider,
|
||||
model: options.model,
|
||||
embedQuery: vi.fn(async () => {
|
||||
if (isOld) {
|
||||
oldEmbeddingStarted();
|
||||
await oldEmbeddingGate;
|
||||
}
|
||||
return isOld ? [0.1] : [0.2];
|
||||
}),
|
||||
embedBatch: vi.fn(async () => [[0.1]]),
|
||||
close: isOld ? closeOldProvider : closeReplacementProvider,
|
||||
},
|
||||
};
|
||||
return providerResult({
|
||||
id: options.provider,
|
||||
model: options.model,
|
||||
embedQuery: vi.fn(async () => {
|
||||
if (isOld) {
|
||||
oldEmbeddingStart.resolve();
|
||||
await oldEmbeddingGate.promise;
|
||||
}
|
||||
return isOld ? [0.1] : [0.2];
|
||||
}),
|
||||
close: isOld ? closeOldProvider : closeReplacementProvider,
|
||||
});
|
||||
},
|
||||
);
|
||||
providerMocks.getMemoryEmbeddingProvider.mockReturnValue({
|
||||
|
|
@ -525,13 +503,13 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
|
||||
try {
|
||||
const oldEmbedding = embed(embeddings, "main", "old config request", oldConfig);
|
||||
await oldEmbeddingStart;
|
||||
await oldEmbeddingStart.promise;
|
||||
const replacementEmbedding = embed(embeddings, "main", "new config request", newConfig);
|
||||
await Promise.resolve();
|
||||
expect(createProvider).toHaveBeenCalledOnce();
|
||||
expect(closeOldProvider).not.toHaveBeenCalled();
|
||||
|
||||
releaseOldEmbedding();
|
||||
oldEmbeddingGate.resolve();
|
||||
await expect(Promise.all([oldEmbedding, replacementEmbedding])).resolves.toEqual([
|
||||
[0.1],
|
||||
[0.2],
|
||||
|
|
@ -557,7 +535,7 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
expectDefined(createProvider.mock.invocationCallOrder[1], "new config create order"),
|
||||
);
|
||||
} finally {
|
||||
releaseOldEmbedding();
|
||||
oldEmbeddingGate.resolve();
|
||||
await embeddings.close?.();
|
||||
}
|
||||
|
||||
|
|
@ -582,14 +560,8 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
},
|
||||
]);
|
||||
};
|
||||
let releaseOldEmbedding: () => void = () => {};
|
||||
const oldEmbeddingGate = new Promise<void>((resolve) => {
|
||||
releaseOldEmbedding = resolve;
|
||||
});
|
||||
let oldEmbeddingStarted: () => void = () => {};
|
||||
const oldEmbeddingStart = new Promise<void>((resolve) => {
|
||||
oldEmbeddingStarted = resolve;
|
||||
});
|
||||
const oldEmbeddingGate = createDeferred<void>();
|
||||
const oldEmbeddingStart = createDeferred<void>();
|
||||
const closeOldProvider = vi.fn(async () => {});
|
||||
const closeReplacementProvider = vi.fn(async () => {});
|
||||
const createProvider = vi.fn(async (options: { agentDir?: string }) => {
|
||||
|
|
@ -602,21 +574,16 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
throw new Error("in-flight agent credential unavailable");
|
||||
}
|
||||
const oldAccount = profile.key === "fixture-inflight-old";
|
||||
return {
|
||||
provider: {
|
||||
id: "openai",
|
||||
model: "text-embedding-3-small",
|
||||
embedQuery: vi.fn(async () => {
|
||||
if (oldAccount) {
|
||||
oldEmbeddingStarted();
|
||||
await oldEmbeddingGate;
|
||||
}
|
||||
return [0.1, 0.2, 0.3];
|
||||
}),
|
||||
embedBatch: vi.fn(async () => [[0.1, 0.2, 0.3]]),
|
||||
close: oldAccount ? closeOldProvider : closeReplacementProvider,
|
||||
},
|
||||
};
|
||||
return providerResult({
|
||||
embedQuery: vi.fn(async () => {
|
||||
if (oldAccount) {
|
||||
oldEmbeddingStart.resolve();
|
||||
await oldEmbeddingGate.promise;
|
||||
}
|
||||
return [0.1, 0.2, 0.3];
|
||||
}),
|
||||
close: oldAccount ? closeOldProvider : closeReplacementProvider,
|
||||
});
|
||||
});
|
||||
providerMocks.getMemoryEmbeddingProvider.mockReturnValue({
|
||||
id: "openai",
|
||||
|
|
@ -634,7 +601,7 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
try {
|
||||
publishCredential("fixture-inflight-old");
|
||||
const firstEmbedding = embed(embeddings, "private", "old account request");
|
||||
await oldEmbeddingStart;
|
||||
await oldEmbeddingStart.promise;
|
||||
|
||||
publishCredential("fixture-inflight-new");
|
||||
const replacementEmbedding = embed(embeddings, "private", "new account request");
|
||||
|
|
@ -642,7 +609,7 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
expect(createProvider).toHaveBeenCalledOnce();
|
||||
expect(closeOldProvider).not.toHaveBeenCalled();
|
||||
|
||||
releaseOldEmbedding();
|
||||
oldEmbeddingGate.resolve();
|
||||
await expect(Promise.all([firstEmbedding, replacementEmbedding])).resolves.toEqual([
|
||||
[0.1, 0.2, 0.3],
|
||||
[0.1, 0.2, 0.3],
|
||||
|
|
@ -655,7 +622,7 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
expectDefined(createProvider.mock.invocationCallOrder[1], "new account create order"),
|
||||
);
|
||||
} finally {
|
||||
releaseOldEmbedding();
|
||||
oldEmbeddingGate.resolve();
|
||||
await embeddings.close?.();
|
||||
clearRuntimeAuthProfileStoreSnapshots();
|
||||
}
|
||||
|
|
@ -664,24 +631,13 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
});
|
||||
|
||||
it("queues replacement behind close intent while provider creation is pending", async () => {
|
||||
let releaseFirstCreate: () => void = () => {};
|
||||
const firstCreateGate = new Promise<void>((resolve) => {
|
||||
releaseFirstCreate = resolve;
|
||||
});
|
||||
const firstCreateGate = createDeferred<void>();
|
||||
const closeProvider = vi.fn(async () => {});
|
||||
const createProvider = vi.fn(async () => {
|
||||
if (createProvider.mock.calls.length === 1) {
|
||||
await firstCreateGate;
|
||||
await firstCreateGate.promise;
|
||||
}
|
||||
return {
|
||||
provider: {
|
||||
id: "openai",
|
||||
model: "text-embedding-3-small",
|
||||
embedQuery: vi.fn(async () => [0.1, 0.2, 0.3]),
|
||||
embedBatch: vi.fn(async () => [[0.1, 0.2, 0.3]]),
|
||||
close: closeProvider,
|
||||
},
|
||||
};
|
||||
return providerResult({ close: closeProvider });
|
||||
});
|
||||
providerMocks.getMemoryEmbeddingProvider.mockReturnValue({
|
||||
id: "openai",
|
||||
|
|
@ -698,7 +654,7 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
await Promise.resolve();
|
||||
expect(createProvider).toHaveBeenCalledTimes(1);
|
||||
|
||||
releaseFirstCreate();
|
||||
firstCreateGate.resolve();
|
||||
await firstEmbed;
|
||||
await closePromise;
|
||||
await replacementEmbed;
|
||||
|
|
@ -722,24 +678,10 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
const closeCurrent = vi.fn(async () => {});
|
||||
const createProvider = vi
|
||||
.fn()
|
||||
.mockResolvedValueOnce({
|
||||
provider: {
|
||||
id: "openai",
|
||||
model: "older",
|
||||
embedQuery: vi.fn(async () => [0.1]),
|
||||
embedBatch: vi.fn(async () => [[0.1]]),
|
||||
close: closeOlder,
|
||||
},
|
||||
})
|
||||
.mockResolvedValueOnce({
|
||||
provider: {
|
||||
id: "openai",
|
||||
model: "current",
|
||||
embedQuery: vi.fn(async () => [0.2]),
|
||||
embedBatch: vi.fn(async () => [[0.2]]),
|
||||
close: closeCurrent,
|
||||
},
|
||||
});
|
||||
.mockResolvedValueOnce(providerResult({ model: "older", vector: [0.1], close: closeOlder }))
|
||||
.mockResolvedValueOnce(
|
||||
providerResult({ model: "current", vector: [0.2], close: closeCurrent }),
|
||||
);
|
||||
providerMocks.getMemoryEmbeddingProvider.mockReturnValue({
|
||||
id: "openai",
|
||||
create: createProvider,
|
||||
|
|
@ -760,35 +702,26 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
});
|
||||
|
||||
it("drains an admitted embedding before provider close", async () => {
|
||||
let markEmbedStarted: () => void = () => {};
|
||||
const embedStarted = new Promise<void>((resolve) => {
|
||||
markEmbedStarted = resolve;
|
||||
});
|
||||
let releaseEmbed: () => void = () => {};
|
||||
const embedGate = new Promise<void>((resolve) => {
|
||||
releaseEmbed = resolve;
|
||||
});
|
||||
const embedStarted = createDeferred<void>();
|
||||
const embedGate = createDeferred<void>();
|
||||
const closeProvider = vi.fn(async () => {});
|
||||
providerMocks.getMemoryEmbeddingProvider.mockReturnValue({
|
||||
id: "openai",
|
||||
create: vi.fn(async () => ({
|
||||
provider: {
|
||||
id: "openai",
|
||||
model: "text-embedding-3-small",
|
||||
create: vi.fn(async () =>
|
||||
providerResult({
|
||||
embedQuery: vi.fn(async () => {
|
||||
markEmbedStarted();
|
||||
await embedGate;
|
||||
embedStarted.resolve();
|
||||
await embedGate.promise;
|
||||
return [0.1, 0.2, 0.3];
|
||||
}),
|
||||
embedBatch: vi.fn(async () => [[0.1, 0.2, 0.3]]),
|
||||
close: closeProvider,
|
||||
},
|
||||
})),
|
||||
}),
|
||||
),
|
||||
});
|
||||
|
||||
const embeddings = createEmbeddings(createApi());
|
||||
const embedPromise = embed(embeddings, "main", "active");
|
||||
await embedStarted;
|
||||
await embedStarted.promise;
|
||||
const closePromise = embeddings.close?.();
|
||||
await Promise.resolve();
|
||||
|
||||
|
|
@ -797,7 +730,7 @@ describe("memory-lancedb provider lifecycle", () => {
|
|||
"memory-lancedb embeddings are closed",
|
||||
);
|
||||
|
||||
releaseEmbed();
|
||||
embedGate.resolve();
|
||||
await expect(embedPromise).resolves.toEqual([0.1, 0.2, 0.3]);
|
||||
await closePromise;
|
||||
expect(closeProvider).toHaveBeenCalledTimes(1);
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ type AgentEmbeddingProvider = {
|
|||
agentDir: string;
|
||||
promise: Promise<MemoryEmbeddingProvider>;
|
||||
activeUses: number;
|
||||
idleWaiters: Set<() => void>;
|
||||
idleResolver?: () => void;
|
||||
};
|
||||
|
||||
type ProviderAdapterLifecycleState = {
|
||||
|
|
@ -219,8 +219,6 @@ class ProviderAdapterEmbeddings implements Embeddings {
|
|||
private unregisterAuthMutationListener: (() => void) | undefined;
|
||||
private closePromise: Promise<void> | null = null;
|
||||
private closed = false;
|
||||
private activeUses = 0;
|
||||
private idleWaiters = new Set<() => void>();
|
||||
|
||||
constructor(private api: OpenClawPluginApi) {}
|
||||
|
||||
|
|
@ -233,7 +231,7 @@ class ProviderAdapterEmbeddings implements Embeddings {
|
|||
}
|
||||
if (existing) {
|
||||
this.providers.delete(agentId);
|
||||
this.retireProvider(existing);
|
||||
void this.retireProviders([existing]).catch(() => undefined);
|
||||
}
|
||||
|
||||
const entry: AgentEmbeddingProvider = {
|
||||
|
|
@ -247,7 +245,6 @@ class ProviderAdapterEmbeddings implements Embeddings {
|
|||
throw err;
|
||||
}),
|
||||
activeUses: 0,
|
||||
idleWaiters: new Set(),
|
||||
};
|
||||
this.providers.set(agentId, entry);
|
||||
return entry;
|
||||
|
|
@ -258,29 +255,22 @@ class ProviderAdapterEmbeddings implements Embeddings {
|
|||
return;
|
||||
}
|
||||
this.embeddingFingerprint = fingerprint;
|
||||
for (const [agentId, entry] of this.providers) {
|
||||
this.providers.delete(agentId);
|
||||
this.retireProvider(entry);
|
||||
}
|
||||
this.retireMatchingProviders(() => true);
|
||||
}
|
||||
|
||||
private retireProvider(entry: AgentEmbeddingProvider): void {
|
||||
const retirement = runProviderAdapterLifecycle(async () => {
|
||||
// Config replacement revokes the old credential immediately, but a request
|
||||
// already admitted under that identity must finish before its client closes.
|
||||
if (entry.activeUses > 0) {
|
||||
await new Promise<void>((resolve) => {
|
||||
entry.idleWaiters.add(resolve);
|
||||
});
|
||||
private retireMatchingProviders(predicate: (entry: AgentEmbeddingProvider) => boolean): void {
|
||||
const entries: AgentEmbeddingProvider[] = [];
|
||||
for (const [agentId, entry] of this.providers) {
|
||||
if (predicate(entry)) {
|
||||
this.providers.delete(agentId);
|
||||
entries.push(entry);
|
||||
}
|
||||
const provider = await entry.promise.catch(() => null);
|
||||
if (provider) {
|
||||
PROVIDER_ADAPTER_LIFECYCLE.retainedProviders.add(provider);
|
||||
}
|
||||
await drainRetainedProviders();
|
||||
});
|
||||
}
|
||||
if (entries.length === 0) {
|
||||
return;
|
||||
}
|
||||
// The next provider create/close retries process-global retained ownership.
|
||||
void retirement.catch(() => undefined);
|
||||
void this.retireProviders(entries).catch(() => undefined);
|
||||
}
|
||||
|
||||
private invalidateProvidersForAuthMutation(event: {
|
||||
|
|
@ -288,43 +278,28 @@ class ProviderAdapterEmbeddings implements Embeddings {
|
|||
affectsInheritedStores: boolean;
|
||||
}): void {
|
||||
const changedAgentDir = event.agentDir ? resolveFilePath(event.agentDir) : undefined;
|
||||
for (const [agentId, entry] of this.providers) {
|
||||
if (!event.affectsInheritedStores && resolveFilePath(entry.agentDir) !== changedAgentDir) {
|
||||
continue;
|
||||
}
|
||||
this.providers.delete(agentId);
|
||||
this.retireProvider(entry);
|
||||
}
|
||||
this.retireMatchingProviders(
|
||||
(entry) =>
|
||||
event.affectsInheritedStores || resolveFilePath(entry.agentDir) === changedAgentDir,
|
||||
);
|
||||
}
|
||||
|
||||
private acquireUse(): () => void {
|
||||
if (this.closed) {
|
||||
throw new Error("memory-lancedb embeddings are closed");
|
||||
}
|
||||
this.activeUses += 1;
|
||||
let released = false;
|
||||
return () => {
|
||||
if (released) {
|
||||
return;
|
||||
}
|
||||
released = true;
|
||||
this.activeUses -= 1;
|
||||
if (this.activeUses === 0) {
|
||||
const waiters = Array.from(this.idleWaiters);
|
||||
this.idleWaiters.clear();
|
||||
for (const resolve of waiters) {
|
||||
resolve();
|
||||
private async retireProviders(entries: AgentEmbeddingProvider[]): Promise<void> {
|
||||
await runProviderAdapterLifecycle(async () => {
|
||||
for (const entry of entries) {
|
||||
// Admission records the entry lease before embed() first yields, so this
|
||||
// covers invalidation, pending creation, and explicit service close.
|
||||
if (entry.activeUses > 0) {
|
||||
await new Promise<void>((resolve) => {
|
||||
entry.idleResolver = resolve;
|
||||
});
|
||||
}
|
||||
const provider = await entry.promise.catch(() => null);
|
||||
if (provider) {
|
||||
PROVIDER_ADAPTER_LIFECYCLE.retainedProviders.add(provider);
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
private async awaitIdle(): Promise<void> {
|
||||
if (this.activeUses === 0) {
|
||||
return;
|
||||
}
|
||||
await new Promise<void>((resolve) => {
|
||||
this.idleWaiters.add(resolve);
|
||||
await drainRetainedProviders();
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -388,44 +363,41 @@ class ProviderAdapterEmbeddings implements Embeddings {
|
|||
embeddingConfig: EmbeddingConfig,
|
||||
timeoutMs?: number,
|
||||
): Promise<number[]> {
|
||||
const releaseUse = this.acquireUse();
|
||||
if (this.closed) {
|
||||
throw new Error("memory-lancedb embeddings are closed");
|
||||
}
|
||||
const embedding = { ...embeddingConfig };
|
||||
const fingerprint = embeddingConfigFingerprint(embedding);
|
||||
this.invalidate(fingerprint);
|
||||
const entry = this.getProvider(normalizeAgentId(agentId), embedding);
|
||||
entry.activeUses += 1;
|
||||
try {
|
||||
const embedding = { ...embeddingConfig };
|
||||
const fingerprint = embeddingConfigFingerprint(embedding);
|
||||
this.invalidate(fingerprint);
|
||||
const entry = this.getProvider(normalizeAgentId(agentId), embedding);
|
||||
entry.activeUses += 1;
|
||||
const provider = await entry.promise;
|
||||
if (!timeoutMs) {
|
||||
return await provider.embedQuery(text);
|
||||
}
|
||||
const controller = new AbortController();
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
try {
|
||||
const provider = await entry.promise;
|
||||
if (!timeoutMs) {
|
||||
return await provider.embedQuery(text);
|
||||
}
|
||||
const controller = new AbortController();
|
||||
let timer: ReturnType<typeof setTimeout> | undefined;
|
||||
try {
|
||||
timer = setTimeout(
|
||||
() => controller.abort(new Error("memory-lancedb embedding timed out")),
|
||||
resolveTimerTimeoutMs(timeoutMs, 1),
|
||||
);
|
||||
timer.unref?.();
|
||||
return await provider.embedQuery(text, { signal: controller.signal });
|
||||
} finally {
|
||||
if (timer) {
|
||||
clearTimeout(timer);
|
||||
}
|
||||
}
|
||||
timer = setTimeout(
|
||||
() => controller.abort(new Error("memory-lancedb embedding timed out")),
|
||||
resolveTimerTimeoutMs(timeoutMs, 1),
|
||||
);
|
||||
timer.unref?.();
|
||||
return await provider.embedQuery(text, { signal: controller.signal });
|
||||
} finally {
|
||||
entry.activeUses -= 1;
|
||||
if (entry.activeUses === 0) {
|
||||
const waiters = Array.from(entry.idleWaiters);
|
||||
entry.idleWaiters.clear();
|
||||
for (const resolve of waiters) {
|
||||
resolve();
|
||||
}
|
||||
if (timer) {
|
||||
clearTimeout(timer);
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
releaseUse();
|
||||
entry.activeUses -= 1;
|
||||
if (entry.activeUses === 0) {
|
||||
// Map removal gives each entry exactly one retirement waiter.
|
||||
const resolveIdle = entry.idleResolver;
|
||||
entry.idleResolver = undefined;
|
||||
resolveIdle?.();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -451,29 +423,11 @@ class ProviderAdapterEmbeddings implements Embeddings {
|
|||
this.closed = true;
|
||||
this.unregisterAuthMutationListener?.();
|
||||
this.unregisterAuthMutationListener = undefined;
|
||||
const providers = Array.from(this.providers.entries());
|
||||
await runProviderAdapterLifecycle(async () => {
|
||||
// Close intent is queued before waiting. Replacement instances therefore remain
|
||||
// behind this owner while already-admitted embeddings drain to completion.
|
||||
await this.awaitIdle();
|
||||
for (const [, entry] of providers) {
|
||||
const provider = await entry.promise.catch(() => null);
|
||||
if (provider) {
|
||||
PROVIDER_ADAPTER_LIFECYCLE.retainedProviders.add(provider);
|
||||
}
|
||||
}
|
||||
try {
|
||||
await drainRetainedProviders();
|
||||
} finally {
|
||||
// Ownership moved to the process-global retained set before draining. Clear the
|
||||
// instance even when another retained provider fails, so successful closes stay final.
|
||||
for (const [agentId, entry] of providers) {
|
||||
if (this.providers.get(agentId) === entry) {
|
||||
this.providers.delete(agentId);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
const providers = Array.from(this.providers.values());
|
||||
this.providers.clear();
|
||||
// Queue close intent before waiting so replacement instances remain behind
|
||||
// every admitted entry and pending provider creation owned by this service.
|
||||
await this.retireProviders(providers);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -608,7 +608,7 @@ describe("memory plugin e2e", () => {
|
|||
]);
|
||||
});
|
||||
|
||||
test("uses provider adapter auth and drains cleanup before service replacement", async () => {
|
||||
test("uses provider adapter auth and propagates service close failures", async () => {
|
||||
const embedQuery = vi.fn(async () => [0.1, 0.2, 0.3]);
|
||||
const closeProvider = vi
|
||||
.fn<() => Promise<void>>()
|
||||
|
|
@ -709,40 +709,12 @@ describe("memory plugin e2e", () => {
|
|||
const service = firstObjectArg(registerService as unknown as MockCallSource, "service");
|
||||
const stop = service.stop as () => Promise<void>;
|
||||
await expect(stop()).rejects.toThrow("provider close failed");
|
||||
|
||||
const replacementRegisterTool = vi.fn();
|
||||
const replacementRegisterService = vi.fn();
|
||||
registerTestPlugin(memoryPlugin, {
|
||||
...mockApi,
|
||||
registerTool: replacementRegisterTool,
|
||||
registerService: replacementRegisterService,
|
||||
});
|
||||
const replacementRecallTool = replacementRegisterTool.mock.calls
|
||||
.map(([tool]) => materializeRegisteredTool(tool))
|
||||
.find((tool) => tool.name === "memory_recall");
|
||||
if (!replacementRecallTool) {
|
||||
throw new Error("expected replacement memory_recall tool registration");
|
||||
}
|
||||
await replacementRecallTool.execute("call-2", { query: "replacement memory" });
|
||||
|
||||
await expect(stop()).resolves.toBeUndefined();
|
||||
expect(closeProvider).toHaveBeenCalledTimes(2);
|
||||
expect(createProvider).toHaveBeenCalledTimes(2);
|
||||
expect(createProvider).toHaveBeenCalledOnce();
|
||||
expect(embedQuery).toHaveBeenCalledWith("project memory", {
|
||||
signal: expect.any(AbortSignal),
|
||||
});
|
||||
expect(
|
||||
expectDefined(closeProvider.mock.invocationCallOrder[1], "retained provider close order"),
|
||||
).toBeLessThan(
|
||||
expectDefined(
|
||||
createProvider.mock.invocationCallOrder[1],
|
||||
"replacement provider create order",
|
||||
),
|
||||
);
|
||||
const replacementService = firstObjectArg(
|
||||
replacementRegisterService as unknown as MockCallSource,
|
||||
"replacement service",
|
||||
);
|
||||
await (replacementService.stop as () => Promise<void>)();
|
||||
} finally {
|
||||
resetMemoryModuleMocks();
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue