@lobehub/lobehub 2.0.5 → 2.0.7

package/src/server/services/memory/userMemory/__tests__/extract.runtime.test.ts

@@ -57,54 +57,123 @@ const createExecutor = (privateOverrides?: Partial<MemoryExtractionPrivateConfig
 };
 
 describe('MemoryExtractionExecutor.resolveRuntimeKeyVaults', () => {
-  it('prefers configured providers/models for gatekeeper, embedding, and layer extractors', () => {
+  it('prefers configured providers/models for gatekeeper, embedding, and layer extractors', async () => {
     const executor = createExecutor({
-      embeddingPreferredProviders: ['provider-e'],
-      agentGateKeeperPreferredModels: ['gate-1'],
-      agentGateKeeperPreferredProviders: ['provider-a', 'provider-b'],
-      agentLayerExtractorPreferredProviders: ['provider-l'],
+      embeddingPreferredProviders: ['provider-c', 'provider-a'],
+      agentGateKeeperPreferredModels: ['model-chat-1', 'vendor-prefix/model-chat-1'],
+      agentGateKeeperPreferredProviders: ['provider-c', 'provider-a'],
+      agentLayerExtractorPreferredProviders: ['provider-c', 'provider-a'],
     });
 
     const runtimeState = createRuntimeState(
       [
-        { abilities: {}, id: 'gate-1', providerId: 'provider-a', type: 'chat' },
-        { abilities: {}, id: 'gate-2', providerId: 'provider-b', type: 'chat' },
-        { abilities: {}, id: 'embed-1', providerId: 'provider-e', type: 'embedding' },
-        { abilities: {}, id: 'layer-ctx', providerId: 'provider-l', type: 'chat' },
-        { abilities: {}, id: 'layer-act', providerId: 'provider-l', type: 'chat' },
-        { abilities: {}, id: 'layer-exp', providerId: 'provider-l', type: 'chat' },
-        { abilities: {}, id: 'layer-id', providerId: 'provider-l', type: 'chat' },
-        { abilities: {}, id: 'layer-pref', providerId: 'provider-l', type: 'chat' },
+        {
+          abilities: {},
+          enabled: true,
+          id: 'model-chat-1',
+          type: 'chat',
+          providerId: 'provider-a',
+        },
+        {
+          abilities: {},
+          enabled: true,
+          id: 'model-embedding-1',
+          type: 'embedding',
+          providerId: 'provider-e',
+        },
+        {
+          abilities: {},
+          enabled: true,
+          id: 'vendor-prefix/model-chat-1',
+          type: 'chat',
+          providerId: 'provider-b',
+        },
+        {
+          abilities: {},
+          enabled: true,
+          id: 'vendor-prefix/model-embedding-1',
+          type: 'embedding',
+          providerId: 'provider-b',
+        },
+        {
+          abilities: {},
+          enabled: false,
+          id: 'model-chat-1',
+          type: 'chat',
+          providerId: 'provider-c',
+        },
+        {
+          abilities: {},
+          enabled: false,
+          id: 'model-embedding-1',
+          type: 'embedding',
+          providerId: 'provider-c',
+        },
       ],
       {
         'provider-a': { apiKey: 'a-key' },
         'provider-b': { apiKey: 'b-key' },
+        'provider-c': { apiKey: 'c-key' },
         'provider-e': { apiKey: 'e-key' },
-        'provider-l': { apiKey: 'l-key' },
       },
     );
 
-    const keyVaults = (executor as any).resolveRuntimeKeyVaults(runtimeState);
+    const keyVaults = await (executor as any).resolveRuntimeKeyVaults(runtimeState);
 
     expect(keyVaults).toMatchObject({
-      'provider-a': { apiKey: 'a-key' }, // gatekeeper picked preferred provider/model
-      'provider-e': { apiKey: 'e-key' }, // embedding honored preferred provider
-      'provider-l': { apiKey: 'l-key' }, // layer extractor models resolved
+      'provider-a': { apiKey: 'a-key' },
+      'provider-e': { apiKey: 'e-key' },
     });
   });
 
-  it('warns and falls back to server provider when no enabled provider satisfies embedding model', () => {
+  it('warns and falls back to server provider when no enabled provider satisfies embedding model', async () => {
     const executor = createExecutor();
     const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
 
     const runtimeState = createRuntimeState(
       [
-        { abilities: {}, id: 'gate-2', providerId: 'provider-b', type: 'chat' },
-        { abilities: {}, id: 'layer-act', providerId: 'provider-l', type: 'chat' },
-        { abilities: {}, id: 'layer-ctx', providerId: 'provider-l', type: 'chat' },
-        { abilities: {}, id: 'layer-exp', providerId: 'provider-l', type: 'chat' },
-        { abilities: {}, id: 'layer-id', providerId: 'provider-l', type: 'chat' },
-        { abilities: {}, id: 'layer-pref', providerId: 'provider-l', type: 'chat' },
+        {
+          abilities: {},
+          enabled: true,
+          id: 'model-chat-1',
+          type: 'chat',
+          providerId: 'provider-a',
+        },
+        {
+          abilities: {},
+          enabled: true,
+          id: 'model-embedding-1',
+          type: 'embedding',
+          providerId: 'provider-e',
+        },
+        {
+          abilities: {},
+          enabled: true,
+          id: 'vendor-prefix/model-chat-1',
+          type: 'chat',
+          providerId: 'provider-b',
+        },
+        {
+          abilities: {},
+          enabled: true,
+          id: 'vendor-prefix/model-embedding-1',
+          type: 'embedding',
+          providerId: 'provider-b',
+        },
+        {
+          abilities: {},
+          enabled: false,
+          id: 'model-chat-1',
+          type: 'chat',
+          providerId: 'provider-c',
+        },
+        {
+          abilities: {},
+          enabled: false,
+          id: 'model-embedding-1',
+          type: 'embedding',
+          providerId: 'provider-c',
+        },
       ],
       {
         'provider-b': { apiKey: 'b-key' },
@@ -112,7 +181,7 @@ describe('MemoryExtractionExecutor.resolveRuntimeKeyVaults', () => {
       },
     );
 
-    const keyVaults = (executor as any).resolveRuntimeKeyVaults(runtimeState);
+    const keyVaults = await (executor as any).resolveRuntimeKeyVaults(runtimeState);
 
     expect(keyVaults).toMatchObject({
       'provider-b': { apiKey: 'b-key' },
@@ -123,4 +192,90 @@ describe('MemoryExtractionExecutor.resolveRuntimeKeyVaults', () => {
 
     warnSpy.mockRestore();
   });
+
+  it('ignores disabled providers when resolving key vaults', async () => {
+    const executor = createExecutor({
+      embeddingPreferredProviders: ['provider-disabled', 'provider-a'],
+    });
+
+    const runtimeState = createRuntimeState(
+      [
+        {
+          abilities: {},
+          enabled: false,
+          id: 'embed-1',
+          type: 'embedding',
+          providerId: 'provider-disabled',
+        },
+        {
+          abilities: {},
+          enabled: true,
+          id: 'embed-1',
+          type: 'embedding',
+          providerId: 'provider-a',
+        },
+      ],
+      {
+        'provider-disabled': { apiKey: 'disabled-key' },
+        'provider-a': { apiKey: 'a-key' },
+      },
+    );
+
+    const keyVaults = await (executor as any).resolveRuntimeKeyVaults(runtimeState);
+
+    expect(keyVaults).toMatchObject({
+      'provider-a': { apiKey: 'a-key' },
+    });
+    expect(keyVaults).not.toHaveProperty('provider-disabled');
+  });
+
+  it('respects preferred provider order when multiple providers have the model', async () => {
+    const executor = createExecutor({
+      agentGateKeeper: {
+        model: 'gate-2',
+        provider: 'provider-a', // fallback provider differs from preferred order
+        apiKey: 'sys-a-key',
+        baseURL: 'https://api-a.example.com',
+        language: 'English',
+      },
+      agentGateKeeperPreferredProviders: ['provider-b', 'provider-a'],
+    });
+
+    const runtimeState = createRuntimeState(
+      [
+        { abilities: {}, enabled: true, id: 'gate-2', type: 'chat', providerId: 'provider-a' },
+        { abilities: {}, enabled: true, id: 'gate-2', type: 'chat', providerId: 'provider-b' },
+      ],
+      {
+        'provider-a': { apiKey: 'a-key' },
+        'provider-b': { apiKey: 'b-key' },
+      },
+    );
+
+    const keyVaults = await (executor as any).resolveRuntimeKeyVaults(runtimeState);
+
+    expect(keyVaults).toMatchObject({
+      'provider-b': { apiKey: 'b-key' }, // picks first preferred provider
+    });
+    expect(keyVaults).not.toHaveProperty('provider-a');
+  });
+
+  it('falls back to configured provider when no enabled models match', async () => {
+    const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+    const executor = createExecutor({
+      agentGateKeeper: { model: 'gate-2', provider: 'provider-fallback', apiKey: 'sys-fb-key' },
+    });
+
+    const runtimeState = createRuntimeState([], {
+      'provider-fallback': { apiKey: 'fb-key' },
+    });
+
+    const keyVaults = await (executor as any).resolveRuntimeKeyVaults(runtimeState);
+
+    expect(keyVaults).toMatchObject({
+      'provider-fallback': { apiKey: 'fb-key' },
+    });
+
+    warnSpy.mockRestore();
+  });
 });

package/src/server/services/memory/userMemory/extract.ts

@@ -296,20 +296,54 @@ const maskSecret = (value?: string) => {
   return `${value.slice(0, 6)}***${value.slice(-4)}`;
 };
 
-const resolveRuntimeAgentConfig = (agent: MemoryAgentConfig, keyVaults?: ProviderKeyVaultMap) => {
-  const provider = agent.provider || 'openai';
-  const { apiKey: userApiKey, baseURL: userBaseURL } = extractCredentialsFromVault(
-    keyVaults?.[normalizeProvider(provider)],
+type ProviderCredential = { apiKey?: string; baseURL?: string };
+
+type RuntimeResolveOptions = {
+  fallback?: ProviderCredential;
+  preferred?: {
+    providerIds?: string[];
+  };
+};
+
+const resolveRuntimeAgentConfig = (
+  agent: MemoryAgentConfig,
+  keyVaults?: ProviderKeyVaultMap,
+  options?: RuntimeResolveOptions,
+) => {
+  const normalizedPreferredProviders = (options?.preferred?.providerIds || [])
+    .map(normalizeProvider)
+    .filter(Boolean);
+
+  const providerOrder = Array.from(
+    new Set([
+      ...normalizedPreferredProviders,
+      normalizeProvider(agent.provider || 'openai'),
+      ...Object.keys(keyVaults || {}),
+    ]),
   );
 
-  // Only use the user baseURL if we are also using their API key; otherwise fall back entirely
-  // to system config to avoid mixing credentials.
-  const useUserCredential = !!userApiKey;
-  const apiKey = useUserCredential ? userApiKey : agent.apiKey;
-  const baseURL = useUserCredential ? userBaseURL || agent.baseURL : agent.baseURL;
-  const source = useUserCredential ? 'user-keyvault' : 'system-config';
+  for (const provider of providerOrder) {
+    const { apiKey: userApiKey, baseURL: userBaseURL } = extractCredentialsFromVault(
+      keyVaults?.[provider],
+    );
+    if (!userApiKey) continue;
+
+    // Only use the user baseURL if we are also using their API key; otherwise fall back entirely
+    // to system config to avoid mixing credentials.
+    return {
+      apiKey: userApiKey,
+      baseURL: userBaseURL || agent.baseURL || options?.fallback?.baseURL,
+      provider,
+      source: 'user-keyvault' as const,
+    };
+  }
 
-  return { apiKey, baseURL, provider, source };
+  return {
+    apiKey: agent.apiKey || options?.fallback?.apiKey,
+    baseURL: agent.baseURL || options?.fallback?.baseURL,
+    provider: agent.provider || 'openai',
+    source: 'system-config' as const,
+  };
 };
 
 const logRuntime = debug('lobe-server:memory:user-memory:runtime');
@@ -329,8 +363,12 @@ const debugRuntimeInit = (
   });
 };
 
-const initRuntimeForAgent = async (agent: MemoryAgentConfig, keyVaults?: ProviderKeyVaultMap) => {
-  const resolved = resolveRuntimeAgentConfig(agent, keyVaults);
+const initRuntimeForAgent = async (
+  agent: MemoryAgentConfig,
+  keyVaults?: ProviderKeyVaultMap,
+  options?: RuntimeResolveOptions,
+) => {
+  const resolved = resolveRuntimeAgentConfig(agent, keyVaults, options);
   debugRuntimeInit(agent, resolved);
 
   if (!resolved.apiKey) {
@@ -1142,7 +1180,7 @@ export class MemoryExtractionExecutor {
             userModel.getUserState(KeyVaultsGateKeeper.getUserKeyVaults),
             this.getAiProviderRuntimeState(job.userId),
           ]);
-          const keyVaults = this.resolveRuntimeKeyVaults(aiProviderRuntimeState);
+          const keyVaults = await this.resolveRuntimeKeyVaults(aiProviderRuntimeState);
           const language = userState.settings?.general?.responseLanguage;
 
           const runtimes = await this.getRuntime(job.userId, keyVaults);
@@ -1827,7 +1865,9 @@ export class MemoryExtractionExecutor {
     return aiInfraRepos.getAiProviderRuntimeState(KeyVaultsGateKeeper.getUserKeyVaults);
   }
 
-  private resolveRuntimeKeyVaults(runtimeState: AiProviderRuntimeState): ProviderKeyVaultMap {
+  private async resolveRuntimeKeyVaults(
+    runtimeState: AiProviderRuntimeState,
+  ): Promise<ProviderKeyVaultMap> {
     const normalizedRuntimeConfig = Object.fromEntries(
       Object.entries(runtimeState.runtimeConfig || {}).map(([providerId, config]) => [
         normalizeProvider(providerId),
@@ -1835,98 +1875,46 @@ export class MemoryExtractionExecutor {
       ]),
     );
 
-    const providerModels = runtimeState.enabledAiModels.reduce<Record<string, Set<string>>>(
-      (acc, model) => {
-        const providerId = normalizeProvider(model.providerId);
-        acc[providerId] = acc[providerId] || new Set<string>();
-        acc[providerId].add(model.id);
-        return acc;
-      },
-      {},
-    );
-
-    const resolveProviderForModel = (
-      modelId: string,
-      fallbackProvider?: string,
-      preferredProviders?: string[],
-      preferredModels?: string[],
-      label?: string,
-    ) => {
-      const providerOrder = Array.from(
-        new Set(
-          [
-            ...(preferredProviders?.map(normalizeProvider) || []),
-            fallbackProvider ? normalizeProvider(fallbackProvider) : undefined,
-            ...Object.keys(providerModels),
-          ].filter(Boolean) as string[],
-        ),
-      );
-
-      const candidateModels = preferredModels && preferredModels.length > 0 ? preferredModels : [];
-
-      for (const providerId of providerOrder) {
-        const models = providerModels[providerId];
-        if (!models) continue;
-        if (models.has(modelId)) return providerId;
-
-        const preferredMatch = candidateModels.find((preferredModel) => models.has(preferredModel));
-        if (preferredMatch) return providerId;
-      }
-      if (fallbackProvider) {
-        console.warn(
-          `[memory-extraction] no enabled provider found for ${label || 'model'} "${modelId}"`,
-          `(preferred ${preferredProviders}), falling back to server-configured provider "${fallbackProvider}".`,
-        );
-
-        return normalizeProvider(fallbackProvider);
-      }
-
-      throw new Error(
-        `Unable to resolve provider for ${label || 'model'} "${modelId}". ` +
-          `Check preferred providers/models configuration.`,
-      );
-    };
-
     const keyVaults: ProviderKeyVaultMap = {};
 
-    const gatekeeperProvider = resolveProviderForModel(
-      this.modelConfig.gateModel,
-      this.privateConfig.agentGateKeeper.provider,
-      this.gatekeeperPreferredProviders,
-      this.gatekeeperPreferredModels,
-      'gatekeeper',
-    );
+    const gatekeeperProvider = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
+      fallbackProvider: this.privateConfig.agentGateKeeper.provider,
+      label: 'gatekeeper',
+      modelId: this.modelConfig.gateModel,
+      preferredModels: this.gatekeeperPreferredModels,
+      preferredProviders: this.gatekeeperPreferredProviders,
+    });
     const gatekeeperRuntime = normalizedRuntimeConfig[gatekeeperProvider];
     if (gatekeeperRuntime?.keyVaults) {
       keyVaults[gatekeeperProvider] = gatekeeperRuntime.keyVaults;
     }
 
-    const embeddingProvider = resolveProviderForModel(
-      this.modelConfig.embeddingsModel,
-      this.privateConfig.embedding.provider,
-      this.embeddingPreferredProviders,
-      this.embeddingPreferredModels,
-      'embedding',
-    );
+    const embeddingProvider = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
+      fallbackProvider: this.privateConfig.embedding.provider,
+      label: 'embedding',
+      modelId: this.modelConfig.embeddingsModel,
+      preferredModels: this.embeddingPreferredModels,
+      preferredProviders: this.embeddingPreferredProviders,
+    });
     const embeddingRuntime = normalizedRuntimeConfig[embeddingProvider];
     if (embeddingRuntime?.keyVaults) {
       keyVaults[embeddingProvider] = embeddingRuntime.keyVaults;
     }
 
-    Object.values(this.modelConfig.layerModels).forEach((model) => {
-      if (!model) return;
-      const providerId = resolveProviderForModel(
-        model,
-        this.privateConfig.agentLayerExtractor.provider,
-        this.layerPreferredProviders,
-        this.layerPreferredModels,
-        'layer extractor',
-      );
+    for (const model of Object.values(this.modelConfig.layerModels)) {
+      if (!model) continue;
+      const providerId = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
+        fallbackProvider: this.privateConfig.agentLayerExtractor.provider,
+        label: 'layer extractor',
+        modelId: model,
+        preferredModels: this.layerPreferredModels,
+        preferredProviders: this.layerPreferredProviders,
+      });
       const runtime = normalizedRuntimeConfig[providerId];
       if (runtime?.keyVaults) {
         keyVaults[providerId] = runtime.keyVaults;
       }
-    });
+    }
 
     return keyVaults;
   }
@@ -1944,10 +1932,46 @@ export class MemoryExtractionExecutor {
     const cached = this.runtimeCache.get(userId);
     if (cached) return cached;
 
+    const embeddingOptions: RuntimeResolveOptions = {
+      fallback: {
+        apiKey: this.privateConfig.embedding.apiKey,
+        baseURL: this.privateConfig.embedding.baseURL,
+      },
+      preferred: { providerIds: this.embeddingPreferredProviders },
+    };
+
+    const gatekeeperOptions: RuntimeResolveOptions = {
+      fallback: {
+        apiKey: this.privateConfig.agentGateKeeper.apiKey,
+        baseURL: this.privateConfig.agentGateKeeper.baseURL,
+      },
+      preferred: { providerIds: this.gatekeeperPreferredProviders },
+    };
+
+    const layerExtractorOptions: RuntimeResolveOptions = {
+      fallback: {
+        apiKey: this.privateConfig.agentLayerExtractor.apiKey,
+        baseURL: this.privateConfig.agentLayerExtractor.baseURL,
+      },
+      preferred: { providerIds: this.layerPreferredProviders },
+    };
+
     const runtimes: RuntimeBundle = {
-      embeddings: await initRuntimeForAgent(this.privateConfig.embedding, keyVaults),
-      gatekeeper: await initRuntimeForAgent(this.privateConfig.agentGateKeeper, keyVaults),
-      layerExtractor: await initRuntimeForAgent(this.privateConfig.agentLayerExtractor, keyVaults),
+      embeddings: await initRuntimeForAgent(
+        { ...this.privateConfig.embedding },
+        keyVaults,
+        embeddingOptions,
+      ),
+      gatekeeper: await initRuntimeForAgent(
+        { ...this.privateConfig.agentGateKeeper },
+        keyVaults,
+        gatekeeperOptions,
+      ),
+      layerExtractor: await initRuntimeForAgent(
+        { ...this.privateConfig.agentLayerExtractor },
+        keyVaults,
+        layerExtractorOptions,
+      ),
     };
 
     this.runtimeCache.set(userId, runtimes);
@@ -1986,7 +2010,7 @@ export class MemoryExtractionExecutor {
             userModel.getUserState(KeyVaultsGateKeeper.getUserKeyVaults),
             this.getAiProviderRuntimeState(params.userId),
           ]);
-          const keyVaults = this.resolveRuntimeKeyVaults(aiProviderRuntimeState);
+          const keyVaults = await this.resolveRuntimeKeyVaults(aiProviderRuntimeState);
           const language = params.language || userState.settings?.general?.responseLanguage;
 
           const runtimes = await this.getRuntime(params.userId, keyVaults);
@@ -2266,7 +2290,7 @@ export class MemoryExtractionWorkflowService {
 
     const url = getWorkflowUrl(WORKFLOW_PATHS.personaUpdate, baseUrl);
     return this.getClient().trigger({
-      body: { userId },
+      body: { userIds: [userId] },
       flowControl: {
         key: `memory-user-memory.pipelines.persona.update-write.${userId}`,
         parallelism: 1,

package/src/server/services/memory/userMemory/persona/__tests__/service.test.ts

@@ -8,6 +8,32 @@ import { UserPersonaModel } from '@/database/models/userMemory/persona';
 
 import { UserPersonaService } from '../service';
 
+// Use var to avoid TDZ with vi.mock hoisting
+var aiInfraMocks:
+  | undefined
+  | {
+      getAiProviderRuntimeState: ReturnType<typeof vi.fn>;
+      tryMatchingModelFrom: ReturnType<typeof vi.fn>;
+      tryMatchingProviderFrom: ReturnType<typeof vi.fn>;
+    };
+
+vi.mock('@/database/repositories/aiInfra', () => {
+  aiInfraMocks = {
+    getAiProviderRuntimeState: vi.fn(),
+    tryMatchingModelFrom: vi.fn(),
+    tryMatchingProviderFrom: vi.fn(),
+  };
+
+  const AiInfraRepos = vi.fn().mockImplementation(() => ({
+    getAiProviderRuntimeState: aiInfraMocks!.getAiProviderRuntimeState,
+  })) as unknown as typeof import('@/database/repositories/aiInfra').AiInfraRepos;
+
+  (AiInfraRepos as any).tryMatchingModelFrom = aiInfraMocks!.tryMatchingModelFrom;
+  (AiInfraRepos as any).tryMatchingProviderFrom = aiInfraMocks!.tryMatchingProviderFrom;
+
+  return { AiInfraRepos };
+});
+
 vi.mock('@/server/globalConfig/parseMemoryExtractionConfig', () => ({
   parseMemoryExtractionConfig: () => ({
     agentLayerExtractor: {
@@ -28,6 +54,10 @@ vi.mock('@/server/globalConfig/parseMemoryExtractionConfig', () => ({
   }),
 }));
 
+vi.mock('@/server/modules/KeyVaultsEncrypt', () => ({
+  KeyVaultsGateKeeper: { getUserKeyVaults: vi.fn() },
+}));
+
 const structuredResult = {
   diff: '- updated',
   memoryIds: ['mem-1'],
@@ -56,6 +86,22 @@ const userId = 'user-persona-service';
 
 beforeEach(async () => {
   toolCall.mockClear();
+  aiInfraMocks!.getAiProviderRuntimeState.mockReset();
+  aiInfraMocks!.tryMatchingModelFrom.mockReset();
+  aiInfraMocks!.tryMatchingProviderFrom.mockReset();
+  aiInfraMocks!.tryMatchingModelFrom.mockResolvedValue('openai');
+  aiInfraMocks!.tryMatchingProviderFrom.mockResolvedValue('openai');
+  aiInfraMocks!.getAiProviderRuntimeState.mockResolvedValue({
+    enabledAiModels: [
+      { abilities: {}, enabled: true, id: 'gpt-mock', providerId: 'openai', type: 'chat' },
+    ],
+    enabledAiProviders: [],
+    enabledChatAiProviders: [],
+    enabledImageAiProviders: [],
+    runtimeConfig: {
+      openai: { keyVaults: { apiKey: 'vault-key', baseURL: 'https://vault.example.com' } },
+    },
+  });
   db = await getTestDB();
 
   await db.delete(users);