@lobehub/lobehub 2.0.6 → 2.0.8

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/CHANGELOG.md CHANGED
@@ -2,6 +2,56 @@
2
2
 
3
3
  # Changelog
4
4
 
5
+ ### [Version 2.0.8](https://github.com/lobehub/lobe-chat/compare/v2.0.7...v2.0.8)
6
+
7
+ <sup>Released on **2026-01-28**</sup>
8
+
9
+ #### 🐛 Bug Fixes
10
+
11
+ - **misc**: Fix inbox agent in mobile.
12
+
13
+ <br/>
14
+
15
+ <details>
16
+ <summary><kbd>Improvements and Fixes</kbd></summary>
17
+
18
+ #### What's fixed
19
+
20
+ - **misc**: Fix inbox agent in mobile, closes [#11929](https://github.com/lobehub/lobe-chat/issues/11929) ([42f5c0b](https://github.com/lobehub/lobe-chat/commit/42f5c0b))
21
+
22
+ </details>
23
+
24
+ <div align="right">
25
+
26
+ [![](https://img.shields.io/badge/-BACK_TO_TOP-151515?style=flat-square)](#readme-top)
27
+
28
+ </div>
29
+
30
+ ### [Version 2.0.7](https://github.com/lobehub/lobe-chat/compare/v2.0.6...v2.0.7)
31
+
32
+ <sup>Released on **2026-01-28**</sup>
33
+
34
+ #### 🐛 Bug Fixes
35
+
36
+ - **model-runtime**: Include tool_calls in speed metrics & add getActiveTraceId.
37
+
38
+ <br/>
39
+
40
+ <details>
41
+ <summary><kbd>Improvements and Fixes</kbd></summary>
42
+
43
+ #### What's fixed
44
+
45
+ - **model-runtime**: Include tool_calls in speed metrics & add getActiveTraceId, closes [#11927](https://github.com/lobehub/lobe-chat/issues/11927) ([b24da44](https://github.com/lobehub/lobe-chat/commit/b24da44))
46
+
47
+ </details>
48
+
49
+ <div align="right">
50
+
51
+ [![](https://img.shields.io/badge/-BACK_TO_TOP-151515?style=flat-square)](#readme-top)
52
+
53
+ </div>
54
+
5
55
  ### [Version 2.0.6](https://github.com/lobehub/lobe-chat/compare/v2.0.5...v2.0.6)
6
56
 
7
57
  <sup>Released on **2026-01-27**</sup>
package/changelog/v2.json CHANGED
@@ -1,4 +1,18 @@
1
1
  [
2
+ {
3
+ "children": {
4
+ "fixes": [
5
+ "Fix inbox agent in mobile."
6
+ ]
7
+ },
8
+ "date": "2026-01-28",
9
+ "version": "2.0.8"
10
+ },
11
+ {
12
+ "children": {},
13
+ "date": "2026-01-28",
14
+ "version": "2.0.7"
15
+ },
2
16
  {
3
17
  "children": {
4
18
  "fixes": [
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@lobehub/lobehub",
3
- "version": "2.0.6",
3
+ "version": "2.0.8",
4
4
  "description": "LobeHub - an open-source,comprehensive AI Agent framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.",
5
5
  "keywords": [
6
6
  "framework",
@@ -2,6 +2,7 @@ import type {
2
2
  AiProviderDetailItem,
3
3
  AiProviderListItem,
4
4
  AiProviderRuntimeConfig,
5
+ AiProviderRuntimeState,
5
6
  EnabledProvider,
6
7
  } from '@lobechat/types';
7
8
  import { AiProviderModelListItem, EnabledAiModel, ExtendParamsType } from 'model-bank';
@@ -1774,4 +1775,55 @@ describe('AiInfraRepos', () => {
1774
1775
  });
1775
1776
  });
1776
1777
  });
1778
+
1779
+ describe('AiInfraRepos.tryMatchingProviderFrom', () => {
1780
+ const createRuntimeState = (models: EnabledAiModel[]): AiProviderRuntimeState => ({
1781
+ enabledAiModels: models,
1782
+ enabledAiProviders: [],
1783
+ enabledChatAiProviders: [],
1784
+ enabledImageAiProviders: [],
1785
+ runtimeConfig: {},
1786
+ });
1787
+
1788
+ it('prefers provider order when multiple providers have model', async () => {
1789
+ const runtimeState = createRuntimeState([
1790
+ { abilities: {}, enabled: true, id: 'm-1', type: 'chat', providerId: 'provider-b' },
1791
+ { abilities: {}, enabled: true, id: 'm-1', type: 'chat', providerId: 'provider-a' },
1792
+ ]);
1793
+
1794
+ const providerId = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
1795
+ modelId: 'm-1',
1796
+ preferredProviders: ['provider-b', 'provider-a'],
1797
+ });
1798
+
1799
+ expect(providerId).toBe('provider-b');
1800
+ });
1801
+
1802
+ it('ignores disabled models when matching', async () => {
1803
+ const runtimeState = createRuntimeState([
1804
+ { abilities: {}, enabled: false, id: 'm-1', type: 'chat', providerId: 'provider-disabled' },
1805
+ { abilities: {}, enabled: true, id: 'm-1', type: 'chat', providerId: 'provider-a' },
1806
+ ]);
1807
+
1808
+ const providerId = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
1809
+ modelId: 'm-1',
1810
+ preferredProviders: ['provider-disabled', 'provider-a'],
1811
+ });
1812
+
1813
+ expect(providerId).toBe('provider-a');
1814
+ });
1815
+
1816
+ it('falls back to provided fallback provider when no match', async () => {
1817
+ const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
1818
+ const runtimeState = createRuntimeState([]);
1819
+
1820
+ const providerId = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
1821
+ modelId: 'm-1',
1822
+ fallbackProvider: 'provider-fallback',
1823
+ });
1824
+
1825
+ expect(providerId).toBe('provider-fallback');
1826
+ warnSpy.mockRestore();
1827
+ });
1828
+ });
1777
1829
  });
@@ -24,6 +24,8 @@ import { LobeChatDatabase } from '../../type';
24
24
 
25
25
  type DecryptUserKeyVaults = (encryptKeyVaultsStr: string | null) => Promise<any>;
26
26
 
27
+ const normalizeProvider = (provider: string) => provider.toLowerCase();
28
+
27
29
  /**
28
30
  * Provider-level search defaults (only used when built-in models don't provide settings.searchImpl and settings.searchProvider)
29
31
  * Note: Not stored in DB, only injected during read
@@ -282,6 +284,107 @@ export class AiInfraRepos {
282
284
  };
283
285
  };
284
286
 
287
+ /**
288
+ * Resolve the best provider for a given model.
289
+ *
290
+ * Matching pipeline:
291
+ * 1) Build a map of provider -> enabled model ids (disabled models are ignored).
292
+ * 2) Walk providers in priority order: preferred providers (if any) -> explicit fallback provider -> remaining providers that have enabled models.
293
+ * 3) For each provider, look for an exact modelId match or any preferred model alias.
294
+ * 4) If nothing matches, fall back to the configured provider (with a warning) or throw when no fallback exists.
295
+ *
296
+ * Handles:
297
+ * - Preferred provider ordering (case-insensitive).
298
+ * - Preferred model aliases.
299
+ * - Disabled models are skipped.
300
+ * - Missing matches: falls back when possible, otherwise surfaces an error.
301
+ *
302
+ * Edge cases to note:
303
+ * - If preferredProviders are set, non-preferred providers are skipped unless they are also the explicit fallback.
304
+ * - If fallbackProvider lacks enabled models, it is still returned (caller should ensure runtimeConfig has credentials).
305
+ */
306
+ static async tryMatchingProviderFrom(
307
+ runtimeState: AiProviderRuntimeState,
308
+ options: {
309
+ fallbackProvider?: string;
310
+ label?: string;
311
+ modelId: string;
312
+ preferredModels?: string[];
313
+ preferredProviders?: string[];
314
+ },
315
+ ): Promise<string> {
316
+ const { modelId, fallbackProvider, preferredModels, preferredProviders, label } = options;
317
+
318
+ // Build a map of provider -> enabled model ids for quick membership checks; skip disabled models entirely
319
+ const providerModels = runtimeState.enabledAiModels.reduce<Record<string, Set<string>>>(
320
+ (acc, model) => {
321
+ if (model.enabled === false) return acc;
322
+
323
+ const providerId = normalizeProvider(model.providerId);
324
+ acc[providerId] = acc[providerId] || new Set<string>();
325
+ acc[providerId].add(model.id);
326
+
327
+ return acc;
328
+ },
329
+ {},
330
+ );
331
+
332
+ // Normalize preferred providers so ordering is stable and comparisons are case-insensitive
333
+ const normalizedPreferredProviders = (preferredProviders || [])
334
+ .map(normalizeProvider)
335
+ .filter(Boolean);
336
+
337
+ // Provider search pipeline:
338
+ // 1) iterate preferred providers (if given)
339
+ // 2) fall back to the explicitly configured fallback provider
340
+ // 3) consider any provider that has enabled models
341
+ const providerOrder = Array.from(
342
+ new Set(
343
+ [
344
+ ...normalizedPreferredProviders,
345
+ fallbackProvider ? normalizeProvider(fallbackProvider) : undefined,
346
+ ...Object.keys(providerModels),
347
+ ].filter(Boolean) as string[],
348
+ ),
349
+ );
350
+
351
+ // Candidate models include the requested modelId plus any preferred model aliases
352
+ const modelTargets = new Set([modelId, ...(preferredModels || [])]);
353
+
354
+ for (const providerId of providerOrder) {
355
+ // If preferred providers are specified, skip non-preferred providers unless they are the explicit fallback
356
+ if (
357
+ normalizedPreferredProviders.length > 0 &&
358
+ providerId !== normalizeProvider(fallbackProvider || '') &&
359
+ !normalizedPreferredProviders.includes(providerId)
360
+ ) {
361
+ continue;
362
+ }
363
+
364
+ const models = providerModels[providerId];
365
+ if (!models) {
366
+ continue;
367
+ }
368
+
369
+ // Accept the first provider in order whose enabled models contain either the requested id or any preferred alias
370
+ const match = Array.from(modelTargets).find((target) => models.has(target));
371
+ if (match) {
372
+ return providerId;
373
+ }
374
+ }
375
+
376
+ if (fallbackProvider) {
377
+ console.warn(
378
+ `[ai-infra] no enabled provider found for ${label || 'model'} "${modelId}" (preferred ${preferredProviders}), falling back to server-configured provider "${fallbackProvider}".`,
379
+ );
380
+ return normalizeProvider(fallbackProvider);
381
+ }
382
+
383
+ throw new Error(
384
+ `Unable to resolve provider for ${label || 'model'} "${modelId}". Check preferred providers/models configuration.`,
385
+ );
386
+ }
387
+
285
388
  getAiProviderModelList = async (
286
389
  providerId: string,
287
390
  options?: {
@@ -472,12 +472,14 @@ export const createTokenSpeedCalculator = (
472
472
  // - text/reasoning: standard text output events
473
473
  // - content_part/reasoning_part: multimodal output events used by Gemini 3+ models
474
474
  // which emit structured parts instead of plain text events
475
+ // - tool_calls: function calling output events
475
476
  if (
476
477
  !outputStartAt &&
477
478
  (chunk.type === 'text' ||
478
479
  chunk.type === 'reasoning' ||
479
480
  chunk.type === 'content_part' ||
480
- chunk.type === 'reasoning_part')
481
+ chunk.type === 'reasoning_part' ||
482
+ chunk.type === 'tool_calls')
481
483
  ) {
482
484
  outputStartAt = Date.now();
483
485
  }
@@ -18,6 +18,7 @@ export { LobeComfyUI } from './providers/comfyui';
18
18
  export { LobeDeepSeekAI } from './providers/deepseek';
19
19
  export { LobeGoogleAI } from './providers/google';
20
20
  export { LobeGroq } from './providers/groq';
21
+ export { LobeHubAI } from './providers/lobehub';
21
22
  export { LobeMinimaxAI } from './providers/minimax';
22
23
  export { LobeMistralAI } from './providers/mistral';
23
24
  export { LobeMoonshotAI } from './providers/moonshot';
@@ -36,7 +37,6 @@ export { LobeXiaomiMiMoAI } from './providers/xiaomimimo';
36
37
  export { LobeZenMuxAI } from './providers/zenmux';
37
38
  export { LobeZeroOneAI } from './providers/zeroone';
38
39
  export { LobeZhipuAI } from './providers/zhipu';
39
- export { LobeHubAI } from './providers/lobehub';
40
40
  export * from './types';
41
41
  export * from './types/error';
42
42
  export { consumeStreamUntilDone } from './utils/consumeStream';
@@ -88,7 +88,6 @@ export const providerRuntimeMap = {
88
88
  deepseek: LobeDeepSeekAI,
89
89
  fal: LobeFalAI,
90
90
  fireworksai: LobeFireworksAI,
91
- lobehub: LobeHubAI,
92
91
  giteeai: LobeGiteeAI,
93
92
  github: LobeGithubAI,
94
93
  google: LobeGoogleAI,
@@ -100,6 +99,7 @@ export const providerRuntimeMap = {
100
99
  internlm: LobeInternLMAI,
101
100
  jina: LobeJinaAI,
102
101
  lmstudio: LobeLMStudioAI,
102
+ lobehub: LobeHubAI,
103
103
  minimax: LobeMinimaxAI,
104
104
  mistral: LobeMistralAI,
105
105
  modelscope: LobeModelScopeAI,
@@ -2,45 +2,35 @@ import { memo } from 'react';
2
2
  import { Link } from 'react-router-dom';
3
3
 
4
4
  import { DEFAULT_INBOX_AVATAR } from '@/const/meta';
5
- import { INBOX_SESSION_ID } from '@/const/session';
6
5
  import { SESSION_CHAT_URL } from '@/const/url';
7
6
  import { useNavigateToAgent } from '@/hooks/useNavigateToAgent';
8
- import { getChatStoreState, useChatStore } from '@/store/chat';
9
- import { chatSelectors } from '@/store/chat/selectors';
7
+ import { useAgentStore } from '@/store/agent';
8
+ import { builtinAgentSelectors } from '@/store/agent/selectors';
10
9
  import { useServerConfigStore } from '@/store/serverConfig';
11
10
  import { useSessionStore } from '@/store/session';
11
+ import { sessionSelectors } from '@/store/session/selectors';
12
12
 
13
13
  import ListItem from '../ListItem';
14
14
 
15
15
  const Inbox = memo(() => {
16
16
  const mobile = useServerConfigStore((s) => s.isMobile);
17
- const activeId = useSessionStore((s) => s.activeId);
17
+ const isInboxActive = useSessionStore(sessionSelectors.isInboxSession);
18
18
  const navigateToAgent = useNavigateToAgent();
19
-
20
- const openNewTopicOrSaveTopic = useChatStore((s) => s.openNewTopicOrSaveTopic);
19
+ const inboxAgentId = useAgentStore(builtinAgentSelectors.inboxAgentId);
21
20
 
22
21
  return (
23
22
  <Link
24
23
  aria-label={'Lobe AI'}
25
- onClick={async (e) => {
24
+ onClick={(e) => {
26
25
  e.preventDefault();
27
- if (activeId === INBOX_SESSION_ID && !mobile) {
28
- // If user tap the inbox again, open a new topic.
29
- // Only for desktop.
30
- const inboxMessages = chatSelectors.inboxActiveTopicMessages(getChatStoreState());
31
- if (inboxMessages.length > 0) {
32
- await openNewTopicOrSaveTopic();
33
- }
34
- } else {
35
- navigateToAgent(INBOX_SESSION_ID);
36
- }
26
+ navigateToAgent(inboxAgentId);
37
27
  }}
38
- to={SESSION_CHAT_URL(INBOX_SESSION_ID, mobile)}
28
+ to={SESSION_CHAT_URL(inboxAgentId, mobile)}
39
29
  >
40
30
  <ListItem
41
- active={activeId === INBOX_SESSION_ID}
31
+ active={isInboxActive}
42
32
  avatar={DEFAULT_INBOX_AVATAR}
43
- key={INBOX_SESSION_ID}
33
+ key={'inbox'}
44
34
  styles={{
45
35
  container: {
46
36
  gap: 12,
@@ -73,9 +73,13 @@ export const useWebUserStateRedirect = () =>
73
73
  }
74
74
 
75
75
  // Redirect away from invite-code page if no longer required
76
+ // Skip redirect if force=true is present (for re-entering invite code)
76
77
  if (pathname.startsWith('/invite-code')) {
77
- window.location.href = '/';
78
- return;
78
+ const params = new URLSearchParams(window.location.search);
79
+ if (params.get('force') !== 'true') {
80
+ window.location.href = '/';
81
+ return;
82
+ }
79
83
  }
80
84
 
81
85
  if (!onboardingSelectors.needsOnboarding(state)) return;
@@ -1,22 +1,23 @@
1
1
  import type { Mock } from 'vitest';
2
2
  import { afterEach, describe, expect, it, vi } from 'vitest';
3
3
 
4
+ // eslint-disable-next-line import/first
5
+ import { getActiveTraceId, injectSpanTraceHeaders } from './traceparent';
6
+
4
7
  vi.mock('@lobechat/observability-otel/api', () => {
5
8
  const inject = vi.fn();
6
9
  const setSpan = vi.fn((_ctx, span) => span);
10
+ const getActiveSpan = vi.fn();
7
11
 
8
12
  return {
9
13
  context: {
10
14
  active: vi.fn(() => ({})),
11
15
  },
12
16
  propagation: { inject },
13
- trace: { setSpan },
17
+ trace: { getActiveSpan, setSpan },
14
18
  };
15
19
  });
16
20
 
17
- // eslint-disable-next-line import/first
18
- import { injectSpanTraceHeaders } from './traceparent';
19
-
20
21
  const mockSpan = (traceId: string, spanId: string) =>
21
22
  ({
22
23
  spanContext: () => ({
@@ -39,7 +40,9 @@ describe('injectSpanTraceHeaders', () => {
39
40
 
40
41
  it('uses propagator output when available', async () => {
41
42
  const { propagation } = await api;
42
- (propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>).mockImplementation((_ctx, carrier) => {
43
+ (
44
+ propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>
45
+ ).mockImplementation((_ctx, carrier) => {
43
46
  carrier.traceparent = 'from-propagator';
44
47
  carrier.tracestate = 'state';
45
48
  });
@@ -56,7 +59,9 @@ describe('injectSpanTraceHeaders', () => {
56
59
 
57
60
  it('falls back to manual traceparent formatting when propagator gives none', async () => {
58
61
  const { propagation } = await api;
59
- (propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>).mockImplementation(() => undefined);
62
+ (
63
+ propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>
64
+ ).mockImplementation(() => undefined);
60
65
 
61
66
  const headers = headersWith();
62
67
  const span = mockSpan('1'.repeat(32), '2'.repeat(16));
@@ -64,6 +69,40 @@ describe('injectSpanTraceHeaders', () => {
64
69
  const tp = injectSpanTraceHeaders(headers, span);
65
70
 
66
71
  expect(tp).toBe('00-11111111111111111111111111111111-2222222222222222-01');
67
- expect(headers.get('traceparent')).toBe('00-11111111111111111111111111111111-2222222222222222-01');
72
+ expect(headers.get('traceparent')).toBe(
73
+ '00-11111111111111111111111111111111-2222222222222222-01',
74
+ );
75
+ });
76
+ });
77
+
78
+ describe('getActiveTraceId', () => {
79
+ const api = vi.importMock<typeof import('@lobechat/observability-otel/api')>(
80
+ '@lobechat/observability-otel/api',
81
+ );
82
+
83
+ afterEach(() => {
84
+ vi.resetAllMocks();
85
+ });
86
+
87
+ it('returns traceId from active span', async () => {
88
+ const { trace } = await api;
89
+ const expectedTraceId = 'a'.repeat(32);
90
+ (trace.getActiveSpan as Mock).mockReturnValue(mockSpan(expectedTraceId, 'b'.repeat(16)));
91
+
92
+ expect(getActiveTraceId()).toBe(expectedTraceId);
93
+ });
94
+
95
+ it('returns undefined when no active span', async () => {
96
+ const { trace } = await api;
97
+ (trace.getActiveSpan as Mock).mockReturnValue(undefined);
98
+
99
+ expect(getActiveTraceId()).toBeUndefined();
100
+ });
101
+
102
+ it('returns undefined when traceId is all zeros', async () => {
103
+ const { trace } = await api;
104
+ (trace.getActiveSpan as Mock).mockReturnValue(mockSpan('0'.repeat(32), 'b'.repeat(16)));
105
+
106
+ expect(getActiveTraceId()).toBeUndefined();
68
107
  });
69
108
  });
@@ -1,13 +1,5 @@
1
- import type {
2
- Span,
3
- Context as OtContext,
4
- TextMapGetter
5
- } from '@lobechat/observability-otel/api';
6
- import {
7
- context as otContext,
8
- propagation,
9
- trace,
10
- } from '@lobechat/observability-otel/api';
1
+ import type { Context as OtContext, Span, TextMapGetter } from '@lobechat/observability-otel/api';
2
+ import { context as otContext, propagation, trace } from '@lobechat/observability-otel/api';
11
3
 
12
4
  // NOTICE: do not try to optimize this into .repeat(...) or similar,
13
5
  // here served for better search / semantic search purpose for further diagnostic
@@ -47,6 +39,16 @@ export const getActiveTraceparent = () => {
47
39
  return toTraceparent(span as Span);
48
40
  };
49
41
 
42
+ /**
43
+ * Get the traceId from the active span.
44
+ */
45
+ export const getActiveTraceId = () => {
46
+ const span = trace.getActiveSpan();
47
+ if (!isValidContext(span)) return undefined;
48
+
49
+ return span!.spanContext().traceId;
50
+ };
51
+
50
52
  /**
51
53
  * Injects the active context into headers using the configured propagator (W3C by default).
52
54
  * Also returns the traceparent for convenience.
@@ -57,54 +57,123 @@ const createExecutor = (privateOverrides?: Partial<MemoryExtractionPrivateConfig
57
57
  };
58
58
 
59
59
  describe('MemoryExtractionExecutor.resolveRuntimeKeyVaults', () => {
60
- it('prefers configured providers/models for gatekeeper, embedding, and layer extractors', () => {
60
+ it('prefers configured providers/models for gatekeeper, embedding, and layer extractors', async () => {
61
61
  const executor = createExecutor({
62
- embeddingPreferredProviders: ['provider-e'],
63
- agentGateKeeperPreferredModels: ['gate-1'],
64
- agentGateKeeperPreferredProviders: ['provider-a', 'provider-b'],
65
- agentLayerExtractorPreferredProviders: ['provider-l'],
62
+ embeddingPreferredProviders: ['provider-c', 'provider-a'],
63
+ agentGateKeeperPreferredModels: ['model-chat-1', 'vendor-prefix/model-chat-1'],
64
+ agentGateKeeperPreferredProviders: ['provider-c', 'provider-a'],
65
+ agentLayerExtractorPreferredProviders: ['provider-c', 'provider-a'],
66
66
  });
67
67
 
68
68
  const runtimeState = createRuntimeState(
69
69
  [
70
- { abilities: {}, id: 'gate-1', providerId: 'provider-a', type: 'chat' },
71
- { abilities: {}, id: 'gate-2', providerId: 'provider-b', type: 'chat' },
72
- { abilities: {}, id: 'embed-1', providerId: 'provider-e', type: 'embedding' },
73
- { abilities: {}, id: 'layer-ctx', providerId: 'provider-l', type: 'chat' },
74
- { abilities: {}, id: 'layer-act', providerId: 'provider-l', type: 'chat' },
75
- { abilities: {}, id: 'layer-exp', providerId: 'provider-l', type: 'chat' },
76
- { abilities: {}, id: 'layer-id', providerId: 'provider-l', type: 'chat' },
77
- { abilities: {}, id: 'layer-pref', providerId: 'provider-l', type: 'chat' },
70
+ {
71
+ abilities: {},
72
+ enabled: true,
73
+ id: 'model-chat-1',
74
+ type: 'chat',
75
+ providerId: 'provider-a',
76
+ },
77
+ {
78
+ abilities: {},
79
+ enabled: true,
80
+ id: 'model-embedding-1',
81
+ type: 'embedding',
82
+ providerId: 'provider-e',
83
+ },
84
+ {
85
+ abilities: {},
86
+ enabled: true,
87
+ id: 'vendor-prefix/model-chat-1',
88
+ type: 'chat',
89
+ providerId: 'provider-b',
90
+ },
91
+ {
92
+ abilities: {},
93
+ enabled: true,
94
+ id: 'vendor-prefix/model-embedding-1',
95
+ type: 'embedding',
96
+ providerId: 'provider-b',
97
+ },
98
+ {
99
+ abilities: {},
100
+ enabled: false,
101
+ id: 'model-chat-1',
102
+ type: 'chat',
103
+ providerId: 'provider-c',
104
+ },
105
+ {
106
+ abilities: {},
107
+ enabled: false,
108
+ id: 'model-embedding-1',
109
+ type: 'embedding',
110
+ providerId: 'provider-c',
111
+ },
78
112
  ],
79
113
  {
80
114
  'provider-a': { apiKey: 'a-key' },
81
115
  'provider-b': { apiKey: 'b-key' },
116
+ 'provider-c': { apiKey: 'c-key' },
82
117
  'provider-e': { apiKey: 'e-key' },
83
- 'provider-l': { apiKey: 'l-key' },
84
118
  },
85
119
  );
86
120
 
87
- const keyVaults = (executor as any).resolveRuntimeKeyVaults(runtimeState);
121
+ const keyVaults = await (executor as any).resolveRuntimeKeyVaults(runtimeState);
88
122
 
89
123
  expect(keyVaults).toMatchObject({
90
- 'provider-a': { apiKey: 'a-key' }, // gatekeeper picked preferred provider/model
91
- 'provider-e': { apiKey: 'e-key' }, // embedding honored preferred provider
92
- 'provider-l': { apiKey: 'l-key' }, // layer extractor models resolved
124
+ 'provider-a': { apiKey: 'a-key' },
125
+ 'provider-e': { apiKey: 'e-key' },
93
126
  });
94
127
  });
95
128
 
96
- it('warns and falls back to server provider when no enabled provider satisfies embedding model', () => {
129
+ it('warns and falls back to server provider when no enabled provider satisfies embedding model', async () => {
97
130
  const executor = createExecutor();
98
131
  const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
99
132
 
100
133
  const runtimeState = createRuntimeState(
101
134
  [
102
- { abilities: {}, id: 'gate-2', providerId: 'provider-b', type: 'chat' },
103
- { abilities: {}, id: 'layer-act', providerId: 'provider-l', type: 'chat' },
104
- { abilities: {}, id: 'layer-ctx', providerId: 'provider-l', type: 'chat' },
105
- { abilities: {}, id: 'layer-exp', providerId: 'provider-l', type: 'chat' },
106
- { abilities: {}, id: 'layer-id', providerId: 'provider-l', type: 'chat' },
107
- { abilities: {}, id: 'layer-pref', providerId: 'provider-l', type: 'chat' },
135
+ {
136
+ abilities: {},
137
+ enabled: true,
138
+ id: 'model-chat-1',
139
+ type: 'chat',
140
+ providerId: 'provider-a',
141
+ },
142
+ {
143
+ abilities: {},
144
+ enabled: true,
145
+ id: 'model-embedding-1',
146
+ type: 'embedding',
147
+ providerId: 'provider-e',
148
+ },
149
+ {
150
+ abilities: {},
151
+ enabled: true,
152
+ id: 'vendor-prefix/model-chat-1',
153
+ type: 'chat',
154
+ providerId: 'provider-b',
155
+ },
156
+ {
157
+ abilities: {},
158
+ enabled: true,
159
+ id: 'vendor-prefix/model-embedding-1',
160
+ type: 'embedding',
161
+ providerId: 'provider-b',
162
+ },
163
+ {
164
+ abilities: {},
165
+ enabled: false,
166
+ id: 'model-chat-1',
167
+ type: 'chat',
168
+ providerId: 'provider-c',
169
+ },
170
+ {
171
+ abilities: {},
172
+ enabled: false,
173
+ id: 'model-embedding-1',
174
+ type: 'embedding',
175
+ providerId: 'provider-c',
176
+ },
108
177
  ],
109
178
  {
110
179
  'provider-b': { apiKey: 'b-key' },
@@ -112,7 +181,7 @@ describe('MemoryExtractionExecutor.resolveRuntimeKeyVaults', () => {
112
181
  },
113
182
  );
114
183
 
115
- const keyVaults = (executor as any).resolveRuntimeKeyVaults(runtimeState);
184
+ const keyVaults = await (executor as any).resolveRuntimeKeyVaults(runtimeState);
116
185
 
117
186
  expect(keyVaults).toMatchObject({
118
187
  'provider-b': { apiKey: 'b-key' },
@@ -123,4 +192,90 @@ describe('MemoryExtractionExecutor.resolveRuntimeKeyVaults', () => {
123
192
 
124
193
  warnSpy.mockRestore();
125
194
  });
195
+
196
+ it('ignores disabled providers when resolving key vaults', async () => {
197
+ const executor = createExecutor({
198
+ embeddingPreferredProviders: ['provider-disabled', 'provider-a'],
199
+ });
200
+
201
+ const runtimeState = createRuntimeState(
202
+ [
203
+ {
204
+ abilities: {},
205
+ enabled: false,
206
+ id: 'embed-1',
207
+ type: 'embedding',
208
+ providerId: 'provider-disabled',
209
+ },
210
+ {
211
+ abilities: {},
212
+ enabled: true,
213
+ id: 'embed-1',
214
+ type: 'embedding',
215
+ providerId: 'provider-a',
216
+ },
217
+ ],
218
+ {
219
+ 'provider-disabled': { apiKey: 'disabled-key' },
220
+ 'provider-a': { apiKey: 'a-key' },
221
+ },
222
+ );
223
+
224
+ const keyVaults = await (executor as any).resolveRuntimeKeyVaults(runtimeState);
225
+
226
+ expect(keyVaults).toMatchObject({
227
+ 'provider-a': { apiKey: 'a-key' },
228
+ });
229
+ expect(keyVaults).not.toHaveProperty('provider-disabled');
230
+ });
231
+
232
+ it('respects preferred provider order when multiple providers have the model', async () => {
233
+ const executor = createExecutor({
234
+ agentGateKeeper: {
235
+ model: 'gate-2',
236
+ provider: 'provider-a', // fallback provider differs from preferred order
237
+ apiKey: 'sys-a-key',
238
+ baseURL: 'https://api-a.example.com',
239
+ language: 'English',
240
+ },
241
+ agentGateKeeperPreferredProviders: ['provider-b', 'provider-a'],
242
+ });
243
+
244
+ const runtimeState = createRuntimeState(
245
+ [
246
+ { abilities: {}, enabled: true, id: 'gate-2', type: 'chat', providerId: 'provider-a' },
247
+ { abilities: {}, enabled: true, id: 'gate-2', type: 'chat', providerId: 'provider-b' },
248
+ ],
249
+ {
250
+ 'provider-a': { apiKey: 'a-key' },
251
+ 'provider-b': { apiKey: 'b-key' },
252
+ },
253
+ );
254
+
255
+ const keyVaults = await (executor as any).resolveRuntimeKeyVaults(runtimeState);
256
+
257
+ expect(keyVaults).toMatchObject({
258
+ 'provider-b': { apiKey: 'b-key' }, // picks first preferred provider
259
+ });
260
+ expect(keyVaults).not.toHaveProperty('provider-a');
261
+ });
262
+
263
+ it('falls back to configured provider when no enabled models match', async () => {
264
+ const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
265
+ const executor = createExecutor({
266
+ agentGateKeeper: { model: 'gate-2', provider: 'provider-fallback', apiKey: 'sys-fb-key' },
267
+ });
268
+
269
+ const runtimeState = createRuntimeState([], {
270
+ 'provider-fallback': { apiKey: 'fb-key' },
271
+ });
272
+
273
+ const keyVaults = await (executor as any).resolveRuntimeKeyVaults(runtimeState);
274
+
275
+ expect(keyVaults).toMatchObject({
276
+ 'provider-fallback': { apiKey: 'fb-key' },
277
+ });
278
+
279
+ warnSpy.mockRestore();
280
+ });
126
281
  });
@@ -227,7 +227,7 @@ export interface TopicBatchWorkflowPayload extends MemoryExtractionPayloadInput
227
227
  userId: string;
228
228
  }
229
229
 
230
- type ProviderKeyVaultMap = Record<
230
+ export type ProviderKeyVaultMap = Record<
231
231
  string,
232
232
  AiProviderRuntimeState['runtimeConfig'][string]['keyVaults'] | undefined
233
233
  >;
@@ -296,27 +296,90 @@ const maskSecret = (value?: string) => {
296
296
  return `${value.slice(0, 6)}***${value.slice(-4)}`;
297
297
  };
298
298
 
299
- const resolveRuntimeAgentConfig = (agent: MemoryAgentConfig, keyVaults?: ProviderKeyVaultMap) => {
300
- const provider = agent.provider || 'openai';
301
- const { apiKey: userApiKey, baseURL: userBaseURL } = extractCredentialsFromVault(
302
- keyVaults?.[normalizeProvider(provider)],
299
+ export type ProviderCredential = { apiKey?: string; baseURL?: string };
300
+
301
+ export type RuntimeResolveOptions = {
302
+ fallback?: ProviderCredential;
303
+ preferred?: {
304
+ providerIds?: string[];
305
+ };
306
+ };
307
+
308
+ export const resolveRuntimeAgentConfig = (
309
+ agent: MemoryAgentConfig,
310
+ keyVaults?: ProviderKeyVaultMap,
311
+ options?: RuntimeResolveOptions,
312
+ ) => {
313
+ const normalizedPreferredProviders = (options?.preferred?.providerIds || [])
314
+ .map(normalizeProvider)
315
+ .filter(Boolean);
316
+
317
+ const providerOrder = Array.from(
318
+ new Set([
319
+ ...normalizedPreferredProviders,
320
+ normalizeProvider(agent.provider || 'openai'),
321
+ ...Object.keys(keyVaults || {}),
322
+ ]),
303
323
  );
304
324
 
305
- // Only use the user baseURL if we are also using their API key; otherwise fall back entirely
306
- // to system config to avoid mixing credentials.
307
- const useUserCredential = !!userApiKey;
308
- const apiKey = useUserCredential ? userApiKey : agent.apiKey;
309
- const baseURL = useUserCredential ? userBaseURL || agent.baseURL : agent.baseURL;
310
- const source = useUserCredential ? 'user-keyvault' : 'system-config';
325
+ for (const provider of providerOrder) {
326
+ if (provider === 'lobehub') {
327
+ debugRuntimeInit(agent, {
328
+ provider,
329
+ source: 'user-vault' as const,
330
+ });
331
+
332
+ return ModelRuntime.initializeWithProvider(provider, {});
333
+ }
311
334
 
312
- return { apiKey, baseURL, provider, source };
335
+ const { apiKey: userApiKey, baseURL: userBaseURL } = extractCredentialsFromVault(
336
+ keyVaults?.[provider],
337
+ );
338
+ if (!userApiKey) {
339
+ console.warn(
340
+ `[memory-extraction] skipping provider ${provider} due to missing API key in user vault`,
341
+ );
342
+ continue;
343
+ }
344
+
345
+ debugRuntimeInit(agent, {
346
+ apiKey: userApiKey,
347
+ baseURL: userBaseURL,
348
+ provider,
349
+ source: 'user-vault' as const,
350
+ });
351
+
352
+ // Only use the user baseURL if we are also using their API key; otherwise fall back entirely
353
+ // to system config to avoid mixing credentials.
354
+ return ModelRuntime.initializeWithProvider(provider, {
355
+ apiKey: userApiKey,
356
+ baseURL: userBaseURL,
357
+ });
358
+ }
359
+
360
+ debugRuntimeInit(agent, {
361
+ apiKey: agent.apiKey || options?.fallback?.apiKey,
362
+ baseURL: agent.baseURL || options?.fallback?.baseURL,
363
+ provider: agent.provider || 'openai',
364
+ source: 'system-config' as const,
365
+ });
366
+
367
+ return ModelRuntime.initializeWithProvider(agent.provider || 'openai', {
368
+ apiKey: agent.apiKey || options?.fallback?.apiKey,
369
+ baseURL: agent.baseURL || options?.fallback?.baseURL,
370
+ });
313
371
  };
314
372
 
315
373
  const logRuntime = debug('lobe-server:memory:user-memory:runtime');
316
374
 
317
375
  const debugRuntimeInit = (
318
376
  agent: MemoryAgentConfig,
319
- resolved: ReturnType<typeof resolveRuntimeAgentConfig>,
377
+ resolved: {
378
+ apiKey?: string;
379
+ baseURL?: string;
380
+ provider: string;
381
+ source: 'user-vault' | 'system-config';
382
+ },
320
383
  ) => {
321
384
  if (!logRuntime.enabled) return;
322
385
  logRuntime('init runtime', {
@@ -329,20 +392,6 @@ const debugRuntimeInit = (
329
392
  });
330
393
  };
331
394
 
332
- const initRuntimeForAgent = async (agent: MemoryAgentConfig, keyVaults?: ProviderKeyVaultMap) => {
333
- const resolved = resolveRuntimeAgentConfig(agent, keyVaults);
334
- debugRuntimeInit(agent, resolved);
335
-
336
- if (!resolved.apiKey) {
337
- throw new Error(`Missing API key for ${resolved.provider} memory extraction runtime`);
338
- }
339
-
340
- return ModelRuntime.initializeWithProvider(resolved.provider, {
341
- apiKey: resolved.apiKey,
342
- baseURL: resolved.baseURL,
343
- });
344
- };
345
-
346
395
  const isTopicExtracted = (metadata?: ChatTopicMetadata | null): boolean => {
347
396
  const extractStatus = metadata?.userMemoryExtractStatus;
348
397
  if (extractStatus) return extractStatus === 'completed';
@@ -1142,7 +1191,7 @@ export class MemoryExtractionExecutor {
1142
1191
  userModel.getUserState(KeyVaultsGateKeeper.getUserKeyVaults),
1143
1192
  this.getAiProviderRuntimeState(job.userId),
1144
1193
  ]);
1145
- const keyVaults = this.resolveRuntimeKeyVaults(aiProviderRuntimeState);
1194
+ const keyVaults = await this.resolveRuntimeKeyVaults(aiProviderRuntimeState);
1146
1195
  const language = userState.settings?.general?.responseLanguage;
1147
1196
 
1148
1197
  const runtimes = await this.getRuntime(job.userId, keyVaults);
@@ -1827,7 +1876,9 @@ export class MemoryExtractionExecutor {
1827
1876
  return aiInfraRepos.getAiProviderRuntimeState(KeyVaultsGateKeeper.getUserKeyVaults);
1828
1877
  }
1829
1878
 
1830
- private resolveRuntimeKeyVaults(runtimeState: AiProviderRuntimeState): ProviderKeyVaultMap {
1879
+ private async resolveRuntimeKeyVaults(
1880
+ runtimeState: AiProviderRuntimeState,
1881
+ ): Promise<ProviderKeyVaultMap> {
1831
1882
  const normalizedRuntimeConfig = Object.fromEntries(
1832
1883
  Object.entries(runtimeState.runtimeConfig || {}).map(([providerId, config]) => [
1833
1884
  normalizeProvider(providerId),
@@ -1835,98 +1886,46 @@ export class MemoryExtractionExecutor {
1835
1886
  ]),
1836
1887
  );
1837
1888
 
1838
- const providerModels = runtimeState.enabledAiModels.reduce<Record<string, Set<string>>>(
1839
- (acc, model) => {
1840
- const providerId = normalizeProvider(model.providerId);
1841
- acc[providerId] = acc[providerId] || new Set<string>();
1842
- acc[providerId].add(model.id);
1843
- return acc;
1844
- },
1845
- {},
1846
- );
1847
-
1848
- const resolveProviderForModel = (
1849
- modelId: string,
1850
- fallbackProvider?: string,
1851
- preferredProviders?: string[],
1852
- preferredModels?: string[],
1853
- label?: string,
1854
- ) => {
1855
- const providerOrder = Array.from(
1856
- new Set(
1857
- [
1858
- ...(preferredProviders?.map(normalizeProvider) || []),
1859
- fallbackProvider ? normalizeProvider(fallbackProvider) : undefined,
1860
- ...Object.keys(providerModels),
1861
- ].filter(Boolean) as string[],
1862
- ),
1863
- );
1864
-
1865
- const candidateModels = preferredModels && preferredModels.length > 0 ? preferredModels : [];
1866
-
1867
- for (const providerId of providerOrder) {
1868
- const models = providerModels[providerId];
1869
- if (!models) continue;
1870
- if (models.has(modelId)) return providerId;
1871
-
1872
- const preferredMatch = candidateModels.find((preferredModel) => models.has(preferredModel));
1873
- if (preferredMatch) return providerId;
1874
- }
1875
- if (fallbackProvider) {
1876
- console.warn(
1877
- `[memory-extraction] no enabled provider found for ${label || 'model'} "${modelId}"`,
1878
- `(preferred ${preferredProviders}), falling back to server-configured provider "${fallbackProvider}".`,
1879
- );
1880
-
1881
- return normalizeProvider(fallbackProvider);
1882
- }
1883
-
1884
- throw new Error(
1885
- `Unable to resolve provider for ${label || 'model'} "${modelId}". ` +
1886
- `Check preferred providers/models configuration.`,
1887
- );
1888
- };
1889
-
1890
1889
  const keyVaults: ProviderKeyVaultMap = {};
1891
1890
 
1892
- const gatekeeperProvider = resolveProviderForModel(
1893
- this.modelConfig.gateModel,
1894
- this.privateConfig.agentGateKeeper.provider,
1895
- this.gatekeeperPreferredProviders,
1896
- this.gatekeeperPreferredModels,
1897
- 'gatekeeper',
1898
- );
1891
+ const gatekeeperProvider = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
1892
+ fallbackProvider: this.privateConfig.agentGateKeeper.provider,
1893
+ label: 'gatekeeper',
1894
+ modelId: this.modelConfig.gateModel,
1895
+ preferredModels: this.gatekeeperPreferredModels,
1896
+ preferredProviders: this.gatekeeperPreferredProviders,
1897
+ });
1899
1898
  const gatekeeperRuntime = normalizedRuntimeConfig[gatekeeperProvider];
1900
1899
  if (gatekeeperRuntime?.keyVaults) {
1901
1900
  keyVaults[gatekeeperProvider] = gatekeeperRuntime.keyVaults;
1902
1901
  }
1903
1902
 
1904
- const embeddingProvider = resolveProviderForModel(
1905
- this.modelConfig.embeddingsModel,
1906
- this.privateConfig.embedding.provider,
1907
- this.embeddingPreferredProviders,
1908
- this.embeddingPreferredModels,
1909
- 'embedding',
1910
- );
1903
+ const embeddingProvider = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
1904
+ fallbackProvider: this.privateConfig.embedding.provider,
1905
+ label: 'embedding',
1906
+ modelId: this.modelConfig.embeddingsModel,
1907
+ preferredModels: this.embeddingPreferredModels,
1908
+ preferredProviders: this.embeddingPreferredProviders,
1909
+ });
1911
1910
  const embeddingRuntime = normalizedRuntimeConfig[embeddingProvider];
1912
1911
  if (embeddingRuntime?.keyVaults) {
1913
1912
  keyVaults[embeddingProvider] = embeddingRuntime.keyVaults;
1914
1913
  }
1915
1914
 
1916
- Object.values(this.modelConfig.layerModels).forEach((model) => {
1917
- if (!model) return;
1918
- const providerId = resolveProviderForModel(
1919
- model,
1920
- this.privateConfig.agentLayerExtractor.provider,
1921
- this.layerPreferredProviders,
1922
- this.layerPreferredModels,
1923
- 'layer extractor',
1924
- );
1915
+ for (const model of Object.values(this.modelConfig.layerModels)) {
1916
+ if (!model) continue;
1917
+ const providerId = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
1918
+ fallbackProvider: this.privateConfig.agentLayerExtractor.provider,
1919
+ label: 'layer extractor',
1920
+ modelId: model,
1921
+ preferredModels: this.layerPreferredModels,
1922
+ preferredProviders: this.layerPreferredProviders,
1923
+ });
1925
1924
  const runtime = normalizedRuntimeConfig[providerId];
1926
1925
  if (runtime?.keyVaults) {
1927
1926
  keyVaults[providerId] = runtime.keyVaults;
1928
1927
  }
1929
- });
1928
+ }
1930
1929
 
1931
1930
  return keyVaults;
1932
1931
  }
@@ -1944,10 +1943,46 @@ export class MemoryExtractionExecutor {
1944
1943
  const cached = this.runtimeCache.get(userId);
1945
1944
  if (cached) return cached;
1946
1945
 
1946
+ const embeddingOptions: RuntimeResolveOptions = {
1947
+ fallback: {
1948
+ apiKey: this.privateConfig.embedding.apiKey,
1949
+ baseURL: this.privateConfig.embedding.baseURL,
1950
+ },
1951
+ preferred: { providerIds: this.embeddingPreferredProviders },
1952
+ };
1953
+
1954
+ const gatekeeperOptions: RuntimeResolveOptions = {
1955
+ fallback: {
1956
+ apiKey: this.privateConfig.agentGateKeeper.apiKey,
1957
+ baseURL: this.privateConfig.agentGateKeeper.baseURL,
1958
+ },
1959
+ preferred: { providerIds: this.gatekeeperPreferredProviders },
1960
+ };
1961
+
1962
+ const layerExtractorOptions: RuntimeResolveOptions = {
1963
+ fallback: {
1964
+ apiKey: this.privateConfig.agentLayerExtractor.apiKey,
1965
+ baseURL: this.privateConfig.agentLayerExtractor.baseURL,
1966
+ },
1967
+ preferred: { providerIds: this.layerPreferredProviders },
1968
+ };
1969
+
1947
1970
  const runtimes: RuntimeBundle = {
1948
- embeddings: await initRuntimeForAgent(this.privateConfig.embedding, keyVaults),
1949
- gatekeeper: await initRuntimeForAgent(this.privateConfig.agentGateKeeper, keyVaults),
1950
- layerExtractor: await initRuntimeForAgent(this.privateConfig.agentLayerExtractor, keyVaults),
1971
+ embeddings: await resolveRuntimeAgentConfig(
1972
+ { ...this.privateConfig.embedding },
1973
+ keyVaults,
1974
+ embeddingOptions,
1975
+ ),
1976
+ gatekeeper: await resolveRuntimeAgentConfig(
1977
+ { ...this.privateConfig.agentGateKeeper },
1978
+ keyVaults,
1979
+ gatekeeperOptions,
1980
+ ),
1981
+ layerExtractor: await resolveRuntimeAgentConfig(
1982
+ { ...this.privateConfig.agentLayerExtractor },
1983
+ keyVaults,
1984
+ layerExtractorOptions,
1985
+ ),
1951
1986
  };
1952
1987
 
1953
1988
  this.runtimeCache.set(userId, runtimes);
@@ -1986,7 +2021,7 @@ export class MemoryExtractionExecutor {
1986
2021
  userModel.getUserState(KeyVaultsGateKeeper.getUserKeyVaults),
1987
2022
  this.getAiProviderRuntimeState(params.userId),
1988
2023
  ]);
1989
- const keyVaults = this.resolveRuntimeKeyVaults(aiProviderRuntimeState);
2024
+ const keyVaults = await this.resolveRuntimeKeyVaults(aiProviderRuntimeState);
1990
2025
  const language = params.language || userState.settings?.general?.responseLanguage;
1991
2026
 
1992
2027
  const runtimes = await this.getRuntime(params.userId, keyVaults);
@@ -8,6 +8,32 @@ import { UserPersonaModel } from '@/database/models/userMemory/persona';
8
8
 
9
9
  import { UserPersonaService } from '../service';
10
10
 
11
+ // Use var to avoid TDZ with vi.mock hoisting
12
+ var aiInfraMocks:
13
+ | undefined
14
+ | {
15
+ getAiProviderRuntimeState: ReturnType<typeof vi.fn>;
16
+ tryMatchingModelFrom: ReturnType<typeof vi.fn>;
17
+ tryMatchingProviderFrom: ReturnType<typeof vi.fn>;
18
+ };
19
+
20
+ vi.mock('@/database/repositories/aiInfra', () => {
21
+ aiInfraMocks = {
22
+ getAiProviderRuntimeState: vi.fn(),
23
+ tryMatchingModelFrom: vi.fn(),
24
+ tryMatchingProviderFrom: vi.fn(),
25
+ };
26
+
27
+ const AiInfraRepos = vi.fn().mockImplementation(() => ({
28
+ getAiProviderRuntimeState: aiInfraMocks!.getAiProviderRuntimeState,
29
+ })) as unknown as typeof import('@/database/repositories/aiInfra').AiInfraRepos;
30
+
31
+ (AiInfraRepos as any).tryMatchingModelFrom = aiInfraMocks!.tryMatchingModelFrom;
32
+ (AiInfraRepos as any).tryMatchingProviderFrom = aiInfraMocks!.tryMatchingProviderFrom;
33
+
34
+ return { AiInfraRepos };
35
+ });
36
+
11
37
  vi.mock('@/server/globalConfig/parseMemoryExtractionConfig', () => ({
12
38
  parseMemoryExtractionConfig: () => ({
13
39
  agentLayerExtractor: {
@@ -28,6 +54,10 @@ vi.mock('@/server/globalConfig/parseMemoryExtractionConfig', () => ({
28
54
  }),
29
55
  }));
30
56
 
57
+ vi.mock('@/server/modules/KeyVaultsEncrypt', () => ({
58
+ KeyVaultsGateKeeper: { getUserKeyVaults: vi.fn() },
59
+ }));
60
+
31
61
  const structuredResult = {
32
62
  diff: '- updated',
33
63
  memoryIds: ['mem-1'],
@@ -45,10 +75,8 @@ vi.mock('@lobechat/memory-user-memory', () => ({
45
75
  })),
46
76
  }));
47
77
 
48
- vi.mock('@lobechat/model-runtime', () => ({
49
- ModelRuntime: {
50
- initializeWithProvider: vi.fn().mockResolvedValue({}),
51
- },
78
+ vi.mock('@/server/services/memory/userMemory/extract', () => ({
79
+ resolveRuntimeAgentConfig: vi.fn().mockResolvedValue({}),
52
80
  }));
53
81
 
54
82
  let db: LobeChatDatabase;
@@ -56,6 +84,22 @@ const userId = 'user-persona-service';
56
84
 
57
85
  beforeEach(async () => {
58
86
  toolCall.mockClear();
87
+ aiInfraMocks!.getAiProviderRuntimeState.mockReset();
88
+ aiInfraMocks!.tryMatchingModelFrom.mockReset();
89
+ aiInfraMocks!.tryMatchingProviderFrom.mockReset();
90
+ aiInfraMocks!.tryMatchingModelFrom.mockResolvedValue('openai');
91
+ aiInfraMocks!.tryMatchingProviderFrom.mockResolvedValue('openai');
92
+ aiInfraMocks!.getAiProviderRuntimeState.mockResolvedValue({
93
+ enabledAiModels: [
94
+ { abilities: {}, enabled: true, id: 'gpt-mock', providerId: 'openai', type: 'chat' },
95
+ ],
96
+ enabledAiProviders: [],
97
+ enabledChatAiProviders: [],
98
+ enabledImageAiProviders: [],
99
+ runtimeConfig: {
100
+ openai: { keyVaults: { apiKey: 'vault-key', baseURL: 'https://vault.example.com' } },
101
+ },
102
+ });
59
103
  db = await getTestDB();
60
104
 
61
105
  await db.delete(users);
@@ -9,16 +9,22 @@ import {
9
9
  type UserPersonaExtractionResult,
10
10
  UserPersonaExtractor,
11
11
  } from '@lobechat/memory-user-memory';
12
- import { ModelRuntime } from '@lobechat/model-runtime';
13
12
  import { desc, eq } from 'drizzle-orm';
14
13
 
15
14
  import { UserMemoryModel } from '@/database/models/userMemory';
16
15
  import { UserPersonaModel } from '@/database/models/userMemory/persona';
16
+ import { AiInfraRepos } from '@/database/repositories/aiInfra';
17
17
  import { LobeChatDatabase } from '@/database/type';
18
18
  import {
19
19
  MemoryAgentConfig,
20
20
  parseMemoryExtractionConfig,
21
21
  } from '@/server/globalConfig/parseMemoryExtractionConfig';
22
+ import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
23
+ import {
24
+ type ProviderKeyVaultMap,
25
+ type RuntimeResolveOptions,
26
+ resolveRuntimeAgentConfig,
27
+ } from '@/server/services/memory/userMemory/extract';
22
28
  import { LayersEnum } from '@/types/userMemory';
23
29
  import { trimBasedOnBatchProbe } from '@/utils/chunkers';
24
30
 
@@ -45,7 +51,6 @@ interface UserPersonaAgentResult {
45
51
  export class UserPersonaService {
46
52
  private readonly preferredLanguage?: string;
47
53
  private readonly db: LobeChatDatabase;
48
- private readonly runtime: ModelRuntime;
49
54
  private readonly agentConfig: MemoryAgentConfig;
50
55
 
51
56
  constructor(db: LobeChatDatabase) {
@@ -54,13 +59,35 @@ export class UserPersonaService {
54
59
  this.db = db;
55
60
  this.preferredLanguage = agentPersonaWriter.language;
56
61
  this.agentConfig = agentPersonaWriter;
57
- this.runtime = ModelRuntime.initializeWithProvider(agentPersonaWriter.provider || 'openai', {
58
- apiKey: agentPersonaWriter.apiKey,
59
- baseURL: agentPersonaWriter.baseURL,
60
- });
61
62
  }
62
63
 
63
64
  async composeWriting(payload: UserPersonaAgentPayload): Promise<UserPersonaAgentResult> {
65
+ const aiInfraRepos = new AiInfraRepos(this.db, payload.userId, {});
66
+ const runtimeState = await aiInfraRepos.getAiProviderRuntimeState(
67
+ KeyVaultsGateKeeper.getUserKeyVaults,
68
+ );
69
+ const providerId = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
70
+ fallbackProvider: this.agentConfig.provider,
71
+ label: 'persona writer',
72
+ modelId: this.agentConfig.model,
73
+ });
74
+
75
+ const keyVaults: ProviderKeyVaultMap = Object.entries(runtimeState.runtimeConfig || {}).reduce(
76
+ (acc, [provider, config]) => {
77
+ acc[provider.toLowerCase()] = config?.keyVaults;
78
+ return acc;
79
+ },
80
+ {} as ProviderKeyVaultMap,
81
+ );
82
+
83
+ const runtime = await resolveRuntimeAgentConfig({ ...this.agentConfig }, keyVaults, {
84
+ fallback: {
85
+ apiKey: this.agentConfig.apiKey,
86
+ baseURL: this.agentConfig.baseURL,
87
+ },
88
+ preferred: { providerIds: [providerId] },
89
+ } satisfies RuntimeResolveOptions);
90
+
64
91
  const personaModel = new UserPersonaModel(this.db, payload.userId);
65
92
  const lastDocument = await personaModel.getLatestPersonaDocument();
66
93
  const existingPersonaBaseline = payload.existingPersona ?? lastDocument?.persona;
@@ -68,7 +95,7 @@ export class UserPersonaService {
68
95
  const extractor = new UserPersonaExtractor({
69
96
  agent: 'user-persona',
70
97
  model: this.agentConfig.model,
71
- modelRuntime: this.runtime,
98
+ modelRuntime: runtime,
72
99
  });
73
100
 
74
101
  const agentResult = await extractor.toolCall({