From dc3c55418a6554ae360a552b60b3d3f63ef528df Mon Sep 17 00:00:00 2001 From: Vincent Koc <25068+vincentkoc@users.noreply.github.com> Date: Sat, 11 Jul 2026 02:16:48 -0700 Subject: [PATCH] refactor(memory): scope embedding service acquisition --- .../src/memory/embedding-local-service.ts | 12 ----- .../memory-core/src/memory/embeddings.test.ts | 40 ++++++++++++++-- .../memory-core/src/memory/embeddings.ts | 4 +- .../src/memory/manager-sync-ops.ts | 3 ++ extensions/memory-core/src/memory/manager.ts | 9 ++++ .../memory-core/src/memory/search-manager.ts | 29 ++++++----- .../memory-core/src/runtime-provider.test.ts | 24 +++++++++- .../memory-core/src/runtime-provider.ts | 48 +++++++++++-------- 8 files changed, 116 insertions(+), 53 deletions(-) diff --git a/extensions/memory-core/src/memory/embedding-local-service.ts b/extensions/memory-core/src/memory/embedding-local-service.ts index eb98c18d0d9e..f783d384e8ec 100644 --- a/extensions/memory-core/src/memory/embedding-local-service.ts +++ b/extensions/memory-core/src/memory/embedding-local-service.ts @@ -7,15 +7,3 @@ export type MemoryCoreAcquireLocalService = ( }, signal?: AbortSignal | null, ) => Promise<{ release: () => void } | undefined>; - -let acquireLocalService: MemoryCoreAcquireLocalService | undefined; - -export function configureMemoryCoreEmbeddingLocalService( - acquire: MemoryCoreAcquireLocalService | undefined, -): void { - acquireLocalService = acquire; -} - -export function getMemoryCoreEmbeddingLocalService(): MemoryCoreAcquireLocalService | undefined { - return acquireLocalService; -} diff --git a/extensions/memory-core/src/memory/embeddings.test.ts b/extensions/memory-core/src/memory/embeddings.test.ts index 7e858f967e84..6a648499b2b4 100644 --- a/extensions/memory-core/src/memory/embeddings.test.ts +++ b/extensions/memory-core/src/memory/embeddings.test.ts @@ -3,7 +3,6 @@ import type { OpenClawConfig } from "openclaw/plugin-sdk/config-contracts"; import type { EmbeddingProviderAdapter } from "openclaw/plugin-sdk/embedding-providers"; import type { MemoryEmbeddingProviderAdapter } from "openclaw/plugin-sdk/memory-core-host-engine-embeddings"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { configureMemoryCoreEmbeddingLocalService } from "./embedding-local-service.js"; import { createEmbeddingProvider, resolveEmbeddingProviderFallbackModel } from "./embeddings.js"; const mockEmbeddingRegistry = vi.hoisted(() => ({ @@ -38,7 +37,10 @@ const missingBedrockCredentialsError = new Error( 'No API key found for provider "bedrock". AWS credentials are not available.', ); -function createOptions(provider: string) { +function createOptions( + provider: string, + acquireLocalService = mockEmbeddingRegistry.acquireLocalService, +) { return { config: { plugins: { @@ -59,6 +61,7 @@ function createOptions(provider: string) { provider, fallback: "none", model: "", + acquireLocalService, }; } @@ -102,11 +105,10 @@ function registerMemoryEmbeddingProvider(adapter: MemoryEmbeddingProviderAdapter describe("createEmbeddingProvider", () => { beforeEach(() => { clearMemoryEmbeddingProviders(); - configureMemoryCoreEmbeddingLocalService(mockEmbeddingRegistry.acquireLocalService); + mockEmbeddingRegistry.acquireLocalService.mockReset(); }); afterEach(() => { - configureMemoryCoreEmbeddingLocalService(undefined); clearMemoryEmbeddingProviders(); }); @@ -194,6 +196,36 @@ describe("createEmbeddingProvider", () => { await expect(result.provider?.embedBatch(["doc"])).resolves.toEqual([[3]]); }); + it("keeps concurrent provider creation bound to each caller's local-service hook", async () => { + const observedHooks: unknown[] = []; + registerGenericEmbeddingProvider({ + id: "openai-compatible", + create: async (options) => { + observedHooks.push( + (options as typeof options & { acquireLocalService?: unknown }).acquireLocalService, + ); + await Promise.resolve(); + return { + provider: { + id: "generic", + model: "generic-model", + embed: async () => [1], + embedBatch: async (inputs) => inputs.map(() => [1]), + }, + }; + }, + }); + const firstAcquire = vi.fn(async () => undefined); + const secondAcquire = vi.fn(async () => undefined); + + await Promise.all([ + createEmbeddingProvider(createOptions("openai-compatible", firstAcquire)), + createEmbeddingProvider(createOptions("openai-compatible", secondAcquire)), + ]); + + expect(observedHooks).toEqual([firstAcquire, secondAcquire]); + }); + it("keeps memory-specific providers authoritative during dual registration", async () => { registerGenericEmbeddingProvider({ id: "openai-compatible", diff --git a/extensions/memory-core/src/memory/embeddings.ts b/extensions/memory-core/src/memory/embeddings.ts index b1a2c9bf69ad..95290ae9f826 100644 --- a/extensions/memory-core/src/memory/embeddings.ts +++ b/extensions/memory-core/src/memory/embeddings.ts @@ -13,7 +13,7 @@ import { type MemoryEmbeddingProviderRuntime, } from "openclaw/plugin-sdk/memory-core-host-engine-embeddings"; import { formatErrorMessage } from "../dreaming-shared.js"; -import { getMemoryCoreEmbeddingLocalService } from "./embedding-local-service.js"; +import type { MemoryCoreAcquireLocalService } from "./embedding-local-service.js"; export type EmbeddingProvider = MemoryEmbeddingProvider; export type EmbeddingProviderId = string; @@ -33,6 +33,7 @@ export type EmbeddingProviderResult = { type CreateEmbeddingProviderOptions = MemoryEmbeddingProviderCreateOptions & { provider: EmbeddingProviderRequest; fallback: EmbeddingProviderFallback; + acquireLocalService?: MemoryCoreAcquireLocalService; }; const DEFAULT_MEMORY_EMBEDDING_PROVIDER = "openai"; @@ -239,7 +240,6 @@ async function createWithAdapter( const createOptions = { ...options, model: resolveProviderModel(adapter, options.model), - acquireLocalService: getMemoryCoreEmbeddingLocalService(), }; const result = await adapter.create(createOptions); return { diff --git a/extensions/memory-core/src/memory/manager-sync-ops.ts b/extensions/memory-core/src/memory/manager-sync-ops.ts index 9a213a46af26..990d7b9b6401 100644 --- a/extensions/memory-core/src/memory/manager-sync-ops.ts +++ b/extensions/memory-core/src/memory/manager-sync-ops.ts @@ -50,6 +50,7 @@ import { import { resolveTimerTimeoutMs } from "openclaw/plugin-sdk/number-runtime"; import { normalizeAgentId } from "openclaw/plugin-sdk/routing"; import { normalizeLowercaseStringOrEmpty } from "openclaw/plugin-sdk/string-coerce-runtime"; +import type { MemoryCoreAcquireLocalService } from "./embedding-local-service.js"; import { createEmbeddingProvider, resolveEmbeddingProviderAdapterId, @@ -295,6 +296,7 @@ function createSessionSyncYield(total: number): () => Promise { } export abstract class MemoryManagerSyncOps { + protected abstract readonly acquireLocalService?: MemoryCoreAcquireLocalService; protected abstract readonly cfg: OpenClawConfig; protected abstract readonly agentId: string; protected abstract readonly workspaceDir: string; @@ -2724,6 +2726,7 @@ export abstract class MemoryManagerSyncOps { const fallbackResult = await createEmbeddingProvider({ config: this.cfg, agentDir: resolveAgentDir(this.cfg, this.agentId), + ...(this.acquireLocalService ? { acquireLocalService: this.acquireLocalService } : {}), ...fallbackRequest, }); diff --git a/extensions/memory-core/src/memory/manager.ts b/extensions/memory-core/src/memory/manager.ts index 0ae86a28c57d..6e905b5e5918 100644 --- a/extensions/memory-core/src/memory/manager.ts +++ b/extensions/memory-core/src/memory/manager.ts @@ -29,6 +29,7 @@ import { } from "openclaw/plugin-sdk/memory-core-host-engine-storage"; import { normalizeAgentId } from "openclaw/plugin-sdk/routing"; import { uniqueValues } from "openclaw/plugin-sdk/string-coerce-runtime"; +import type { MemoryCoreAcquireLocalService } from "./embedding-local-service.js"; import { createEmbeddingProvider, resolveEmbeddingProviderAdapterTransport, @@ -294,6 +295,7 @@ async function closeMemoryIndexManagersForScope(params: { export class MemoryIndexManager extends MemoryManagerEmbeddingOps implements MemorySearchManager { private readonly cacheKey: string; private readonly purpose: MemoryIndexManagerPurpose; + protected override readonly acquireLocalService?: MemoryCoreAcquireLocalService; protected readonly cfg: OpenClawConfig; protected readonly agentId: string; protected readonly workspaceDir: string; @@ -372,10 +374,12 @@ export class MemoryIndexManager extends MemoryManagerEmbeddingOps implements Mem cfg: OpenClawConfig; agentId: string; settings: ResolvedMemorySearchConfig; + acquireLocalService?: MemoryCoreAcquireLocalService; }): Promise { return await createEmbeddingProvider({ config: params.cfg, agentDir: resolveAgentDir(params.cfg, params.agentId), + ...(params.acquireLocalService ? { acquireLocalService: params.acquireLocalService } : {}), ...resolveMemoryPrimaryProviderRequest({ settings: params.settings }), }); } @@ -384,6 +388,7 @@ export class MemoryIndexManager extends MemoryManagerEmbeddingOps implements Mem cfg: OpenClawConfig; agentId: string; purpose?: MemoryIndexManagerPurpose; + acquireLocalService?: MemoryCoreAcquireLocalService; }): Promise { const { cfg, agentId } = params; const settings = resolveMemorySearchConfig(cfg, agentId); @@ -428,6 +433,7 @@ export class MemoryIndexManager extends MemoryManagerEmbeddingOps implements Mem settings, providerRequirement, purpose: params.purpose, + acquireLocalService: params.acquireLocalService, }); // Lightweight dirty-file detection for status mode: check for unindexed // session files on disk without triggering a full sync. This runs before @@ -454,10 +460,12 @@ export class MemoryIndexManager extends MemoryManagerEmbeddingOps implements Mem providerRequirement: MemoryEmbeddingProviderRequirement; providerResult?: EmbeddingProviderResult; purpose?: MemoryIndexManagerPurpose; + acquireLocalService?: MemoryCoreAcquireLocalService; }) { super(); const effectiveSettings = resolveEffectiveMemorySearchSettings(params.settings); this.cacheKey = params.cacheKey; + this.acquireLocalService = params.acquireLocalService; this.purpose = params.purpose === "status" || params.purpose === "cli" ? params.purpose : "default"; this.cfg = params.cfg; @@ -550,6 +558,7 @@ export class MemoryIndexManager extends MemoryManagerEmbeddingOps implements Mem cfg: this.cfg, agentId: this.agentId, settings: this.settings, + acquireLocalService: this.acquireLocalService, }); this.applyProviderResult(providerResult); this.providerKey = this.computeProviderKey(); diff --git a/extensions/memory-core/src/memory/search-manager.ts b/extensions/memory-core/src/memory/search-manager.ts index 04c4dab751c1..665bca56fef6 100644 --- a/extensions/memory-core/src/memory/search-manager.ts +++ b/extensions/memory-core/src/memory/search-manager.ts @@ -25,6 +25,7 @@ import { type ResolvedQmdConfig, } from "openclaw/plugin-sdk/memory-core-host-engine-storage"; import { normalizeAgentId } from "openclaw/plugin-sdk/routing"; +import type { MemoryCoreAcquireLocalService } from "./embedding-local-service.js"; const MEMORY_SEARCH_MANAGER_CACHE_KEY = Symbol.for("openclaw.memorySearchManagerCache"); type Maybe = T | null; @@ -125,6 +126,12 @@ export type MemorySearchManagerResult = { }; export type MemorySearchManagerPurpose = "default" | "status" | "cli"; +type MemorySearchManagerParams = { + cfg: OpenClawConfig; + agentId: string; + purpose?: MemorySearchManagerPurpose; + acquireLocalService?: MemoryCoreAcquireLocalService; +}; function getActiveQmdManagerOpenFailure( scopeKey: string, @@ -182,11 +189,9 @@ function applyManagerDebug( }; } -export async function getMemorySearchManager(params: { - cfg: OpenClawConfig; - agentId: string; - purpose?: MemorySearchManagerPurpose; -}): Promise { +export async function getMemorySearchManager( + params: MemorySearchManagerParams, +): Promise { const acquireStartedAt = Date.now(); const purpose = params.purpose ?? "default"; const finish = ( @@ -414,11 +419,7 @@ export async function getMemorySearchManager(params: { } async function getBuiltinMemorySearchManagerAfterQmdFailure( - params: { - cfg: OpenClawConfig; - agentId: string; - purpose?: MemorySearchManagerPurpose; - }, + params: MemorySearchManagerParams, qmdFailureReason: string | undefined, ): Promise { const fallback = await getBuiltinMemorySearchManager(params); @@ -434,11 +435,9 @@ async function getBuiltinMemorySearchManagerAfterQmdFailure( }; } -async function getBuiltinMemorySearchManager(params: { - cfg: OpenClawConfig; - agentId: string; - purpose?: MemorySearchManagerPurpose; -}): Promise { +async function getBuiltinMemorySearchManager( + params: MemorySearchManagerParams, +): Promise { try { const { MemoryIndexManager } = await loadManagerRuntime(); const manager = await MemoryIndexManager.get(params); diff --git a/extensions/memory-core/src/runtime-provider.test.ts b/extensions/memory-core/src/runtime-provider.test.ts index b625e8aba482..1e02967abe60 100644 --- a/extensions/memory-core/src/runtime-provider.test.ts +++ b/extensions/memory-core/src/runtime-provider.test.ts @@ -24,7 +24,7 @@ vi.mock("./memory/index.js", () => ({ getMemorySearchManager: getMemorySearchManagerMock, })); -import { memoryRuntime } from "./runtime-provider.js"; +import { createMemoryRuntime, memoryRuntime } from "./runtime-provider.js"; describe("memoryRuntime", () => { it("preserves manager debug metadata", async () => { @@ -41,4 +41,26 @@ describe("memoryRuntime", () => { agentId: "main", }); }); + + it("keeps local-service acquisition scoped to each runtime instance", async () => { + const cfg = {} as OpenClawConfig; + const firstAcquire = vi.fn(async () => undefined); + const secondAcquire = vi.fn(async () => undefined); + + await Promise.all([ + createMemoryRuntime(firstAcquire).getMemorySearchManager({ cfg, agentId: "first" }), + createMemoryRuntime(secondAcquire).getMemorySearchManager({ cfg, agentId: "second" }), + ]); + + expect(getMemorySearchManagerMock).toHaveBeenCalledWith({ + cfg, + agentId: "first", + acquireLocalService: firstAcquire, + }); + expect(getMemorySearchManagerMock).toHaveBeenCalledWith({ + cfg, + agentId: "second", + acquireLocalService: secondAcquire, + }); + }); }); diff --git a/extensions/memory-core/src/runtime-provider.ts b/extensions/memory-core/src/runtime-provider.ts index 8c0a2531bfad..219cfcad3ab7 100644 --- a/extensions/memory-core/src/runtime-provider.ts +++ b/extensions/memory-core/src/runtime-provider.ts @@ -1,28 +1,38 @@ // Memory Core provider module implements model/runtime integration. import type { MemoryPluginRuntime } from "openclaw/plugin-sdk/memory-core-host-runtime-core"; import { resolveMemoryBackendConfig } from "openclaw/plugin-sdk/memory-core-host-runtime-files"; +import type { MemoryCoreAcquireLocalService } from "./memory/embedding-local-service.js"; import { closeAllMemorySearchManagers, closeMemorySearchManager, getMemorySearchManager, } from "./memory/index.js"; -export const memoryRuntime: MemoryPluginRuntime = { - async getMemorySearchManager(params) { - const { manager, debug, error } = await getMemorySearchManager(params); - return { - manager, - debug, - error, - }; - }, - resolveMemoryBackendConfig(params) { - return resolveMemoryBackendConfig(params); - }, - async closeAllMemorySearchManagers() { - await closeAllMemorySearchManagers(); - }, - async closeMemorySearchManager(params) { - await closeMemorySearchManager(params); - }, -}; +export function createMemoryRuntime( + acquireLocalService?: MemoryCoreAcquireLocalService, +): MemoryPluginRuntime { + return { + async getMemorySearchManager(params) { + const { manager, debug, error } = await getMemorySearchManager({ + ...params, + ...(acquireLocalService ? { acquireLocalService } : {}), + }); + return { + manager, + debug, + error, + }; + }, + resolveMemoryBackendConfig(params) { + return resolveMemoryBackendConfig(params); + }, + async closeAllMemorySearchManagers() { + await closeAllMemorySearchManagers(); + }, + async closeMemorySearchManager(params) { + await closeMemorySearchManager(params); + }, + }; +} + +export const memoryRuntime = createMemoryRuntime();