@lobehub/lobehub 2.0.6 → 2.0.7

package/CHANGELOG.md

@@ -2,6 +2,31 @@
 
 # Changelog
 
+### [Version 2.0.7](https://github.com/lobehub/lobe-chat/compare/v2.0.6...v2.0.7)
+
+<sup>Released on **2026-01-28**</sup>
+
+#### 🐛 Bug Fixes
+
+- **model-runtime**: Include tool_calls in speed metrics & add getActiveTraceId.
+
+<br/>
+
+<details>
+<summary><kbd>Improvements and Fixes</kbd></summary>
+
+#### What's fixed
+
+- **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))
+
+</details>
+
+<div align="right">
+
+[![](https://img.shields.io/badge/-BACK_TO_TOP-151515?style=flat-square)](#readme-top)
+
+</div>
+
 ### [Version 2.0.6](https://github.com/lobehub/lobe-chat/compare/v2.0.5...v2.0.6)
 
 <sup>Released on **2026-01-27**</sup>

package/changelog/v2.json

@@ -1,4 +1,9 @@
 [
+  {
+    "children": {},
+    "date": "2026-01-28",
+    "version": "2.0.7"
+  },
   {
     "children": {
       "fixes": [

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lobehub/lobehub",
-  "version": "2.0.6",
+  "version": "2.0.7",
   "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.",
   "keywords": [
     "framework",

package/packages/database/src/repositories/aiInfra/index.test.ts

@@ -2,6 +2,7 @@ import type {
   AiProviderDetailItem,
   AiProviderListItem,
   AiProviderRuntimeConfig,
+  AiProviderRuntimeState,
   EnabledProvider,
 } from '@lobechat/types';
 import { AiProviderModelListItem, EnabledAiModel, ExtendParamsType } from 'model-bank';
@@ -1774,4 +1775,55 @@ describe('AiInfraRepos', () => {
       });
     });
   });
+
+  describe('AiInfraRepos.tryMatchingProviderFrom', () => {
+    const createRuntimeState = (models: EnabledAiModel[]): AiProviderRuntimeState => ({
+      enabledAiModels: models,
+      enabledAiProviders: [],
+      enabledChatAiProviders: [],
+      enabledImageAiProviders: [],
+      runtimeConfig: {},
+    });
+
+    it('prefers provider order when multiple providers have model', async () => {
+      const runtimeState = createRuntimeState([
+        { abilities: {}, enabled: true, id: 'm-1', type: 'chat', providerId: 'provider-b' },
+        { abilities: {}, enabled: true, id: 'm-1', type: 'chat', providerId: 'provider-a' },
+      ]);
+
+      const providerId = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
+        modelId: 'm-1',
+        preferredProviders: ['provider-b', 'provider-a'],
+      });
+
+      expect(providerId).toBe('provider-b');
+    });
+
+    it('ignores disabled models when matching', async () => {
+      const runtimeState = createRuntimeState([
+        { abilities: {}, enabled: false, id: 'm-1', type: 'chat', providerId: 'provider-disabled' },
+        { abilities: {}, enabled: true, id: 'm-1', type: 'chat', providerId: 'provider-a' },
+      ]);
+
+      const providerId = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
+        modelId: 'm-1',
+        preferredProviders: ['provider-disabled', 'provider-a'],
+      });
+
+      expect(providerId).toBe('provider-a');
+    });
+
+    it('falls back to provided fallback provider when no match', async () => {
+      const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+      const runtimeState = createRuntimeState([]);
+
+      const providerId = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
+        modelId: 'm-1',
+        fallbackProvider: 'provider-fallback',
+      });
+
+      expect(providerId).toBe('provider-fallback');
+      warnSpy.mockRestore();
+    });
+  });
 });

package/packages/database/src/repositories/aiInfra/index.ts

@@ -24,6 +24,8 @@ import { LobeChatDatabase } from '../../type';
 
 type DecryptUserKeyVaults = (encryptKeyVaultsStr: string | null) => Promise<any>;
 
+const normalizeProvider = (provider: string) => provider.toLowerCase();
+
 /**
  * Provider-level search defaults (only used when built-in models don't provide settings.searchImpl and settings.searchProvider)
  * Note: Not stored in DB, only injected during read
@@ -282,6 +284,107 @@ export class AiInfraRepos {
     };
   };
 
+  /**
+   * Resolve the best provider for a given model.
+   *
+   * Matching pipeline:
+   * 1) Build a map of provider -> enabled model ids (disabled models are ignored).
+   * 2) Walk providers in priority order: preferred providers (if any) -> explicit fallback provider -> remaining providers that have enabled models.
+   * 3) For each provider, look for an exact modelId match or any preferred model alias.
+   * 4) If nothing matches, fall back to the configured provider (with a warning) or throw when no fallback exists.
+   *
+   * Handles:
+   * - Preferred provider ordering (case-insensitive).
+   * - Preferred model aliases.
+   * - Disabled models are skipped.
+   * - Missing matches: falls back when possible, otherwise surfaces an error.
+   *
+   * Edge cases to note:
+   * - If preferredProviders are set, non-preferred providers are skipped unless they are also the explicit fallback.
+   * - If fallbackProvider lacks enabled models, it is still returned (caller should ensure runtimeConfig has credentials).
+   */
+  static async tryMatchingProviderFrom(
+    runtimeState: AiProviderRuntimeState,
+    options: {
+      fallbackProvider?: string;
+      label?: string;
+      modelId: string;
+      preferredModels?: string[];
+      preferredProviders?: string[];
+    },
+  ): Promise<string> {
+    const { modelId, fallbackProvider, preferredModels, preferredProviders, label } = options;
+
+    // Build a map of provider -> enabled model ids for quick membership checks; skip disabled models entirely
+    const providerModels = runtimeState.enabledAiModels.reduce<Record<string, Set<string>>>(
+      (acc, model) => {
+        if (model.enabled === false) return acc;
+
+        const providerId = normalizeProvider(model.providerId);
+        acc[providerId] = acc[providerId] || new Set<string>();
+        acc[providerId].add(model.id);
+
+        return acc;
+      },
+      {},
+    );
+
+    // Normalize preferred providers so ordering is stable and comparisons are case-insensitive
+    const normalizedPreferredProviders = (preferredProviders || [])
+      .map(normalizeProvider)
+      .filter(Boolean);
+
+    // Provider search pipeline:
+    // 1) iterate preferred providers (if given)
+    // 2) fall back to the explicitly configured fallback provider
+    // 3) consider any provider that has enabled models
+    const providerOrder = Array.from(
+      new Set(
+        [
+          ...normalizedPreferredProviders,
+          fallbackProvider ? normalizeProvider(fallbackProvider) : undefined,
+          ...Object.keys(providerModels),
+        ].filter(Boolean) as string[],
+      ),
+    );
+
+    // Candidate models include the requested modelId plus any preferred model aliases
+    const modelTargets = new Set([modelId, ...(preferredModels || [])]);
+
+    for (const providerId of providerOrder) {
+      // If preferred providers are specified, skip non-preferred providers unless they are the explicit fallback
+      if (
+        normalizedPreferredProviders.length > 0 &&
+        providerId !== normalizeProvider(fallbackProvider || '') &&
+        !normalizedPreferredProviders.includes(providerId)
+      ) {
+        continue;
+      }
+
+      const models = providerModels[providerId];
+      if (!models) {
+        continue;
+      }
+
+      // Accept the first provider in order whose enabled models contain either the requested id or any preferred alias
+      const match = Array.from(modelTargets).find((target) => models.has(target));
+      if (match) {
+        return providerId;
+      }
+    }
+
+    if (fallbackProvider) {
+      console.warn(
+        `[ai-infra] 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.`,
+    );
+  }
+
   getAiProviderModelList = async (
     providerId: string,
     options?: {

package/packages/model-runtime/src/core/streams/protocol.ts

@@ -472,12 +472,14 @@ export const createTokenSpeedCalculator = (
     // - text/reasoning: standard text output events
     // - content_part/reasoning_part: multimodal output events used by Gemini 3+ models
     //   which emit structured parts instead of plain text events
+    // - tool_calls: function calling output events
     if (
       !outputStartAt &&
       (chunk.type === 'text' ||
         chunk.type === 'reasoning' ||
         chunk.type === 'content_part' ||
-        chunk.type === 'reasoning_part')
+        chunk.type === 'reasoning_part' ||
+        chunk.type === 'tool_calls')
     ) {
       outputStartAt = Date.now();
     }

package/src/layout/GlobalProvider/useUserStateRedirect.ts

@@ -73,9 +73,13 @@ export const useWebUserStateRedirect = () =>
     }
 
     // Redirect away from invite-code page if no longer required
+    // Skip redirect if force=true is present (for re-entering invite code)
     if (pathname.startsWith('/invite-code')) {
-      window.location.href = '/';
-      return;
+      const params = new URLSearchParams(window.location.search);
+      if (params.get('force') !== 'true') {
+        window.location.href = '/';
+        return;
+      }
     }
 
     if (!onboardingSelectors.needsOnboarding(state)) return;

package/src/libs/observability/traceparent.test.ts

@@ -1,22 +1,23 @@
 import type { Mock } from 'vitest';
 import { afterEach, describe, expect, it, vi } from 'vitest';
 
+// eslint-disable-next-line import/first
+import { getActiveTraceId, injectSpanTraceHeaders } from './traceparent';
+
 vi.mock('@lobechat/observability-otel/api', () => {
   const inject = vi.fn();
   const setSpan = vi.fn((_ctx, span) => span);
+  const getActiveSpan = vi.fn();
 
   return {
     context: {
       active: vi.fn(() => ({})),
     },
     propagation: { inject },
-    trace: { setSpan },
+    trace: { getActiveSpan, setSpan },
   };
 });
 
-// eslint-disable-next-line import/first
-import { injectSpanTraceHeaders } from './traceparent';
-
 const mockSpan = (traceId: string, spanId: string) =>
   ({
     spanContext: () => ({
@@ -39,7 +40,9 @@ describe('injectSpanTraceHeaders', () => {
 
   it('uses propagator output when available', async () => {
     const { propagation } = await api;
-    (propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>).mockImplementation((_ctx, carrier) => {
+    (
+      propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>
+    ).mockImplementation((_ctx, carrier) => {
       carrier.traceparent = 'from-propagator';
       carrier.tracestate = 'state';
     });
@@ -56,7 +59,9 @@ describe('injectSpanTraceHeaders', () => {
 
   it('falls back to manual traceparent formatting when propagator gives none', async () => {
     const { propagation } = await api;
-    (propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>).mockImplementation(() => undefined);
+    (
+      propagation.inject as unknown as Mock<typeof propagation.inject<Record<string, string>>>
+    ).mockImplementation(() => undefined);
 
     const headers = headersWith();
     const span = mockSpan('1'.repeat(32), '2'.repeat(16));
@@ -64,6 +69,40 @@ describe('injectSpanTraceHeaders', () => {
     const tp = injectSpanTraceHeaders(headers, span);
 
     expect(tp).toBe('00-11111111111111111111111111111111-2222222222222222-01');
-    expect(headers.get('traceparent')).toBe('00-11111111111111111111111111111111-2222222222222222-01');
+    expect(headers.get('traceparent')).toBe(
+      '00-11111111111111111111111111111111-2222222222222222-01',
+    );
+  });
+});
+
+describe('getActiveTraceId', () => {
+  const api = vi.importMock<typeof import('@lobechat/observability-otel/api')>(
+    '@lobechat/observability-otel/api',
+  );
+
+  afterEach(() => {
+    vi.resetAllMocks();
+  });
+
+  it('returns traceId from active span', async () => {
+    const { trace } = await api;
+    const expectedTraceId = 'a'.repeat(32);
+    (trace.getActiveSpan as Mock).mockReturnValue(mockSpan(expectedTraceId, 'b'.repeat(16)));
+
+    expect(getActiveTraceId()).toBe(expectedTraceId);
+  });
+
+  it('returns undefined when no active span', async () => {
+    const { trace } = await api;
+    (trace.getActiveSpan as Mock).mockReturnValue(undefined);
+
+    expect(getActiveTraceId()).toBeUndefined();
+  });
+
+  it('returns undefined when traceId is all zeros', async () => {
+    const { trace } = await api;
+    (trace.getActiveSpan as Mock).mockReturnValue(mockSpan('0'.repeat(32), 'b'.repeat(16)));
+
+    expect(getActiveTraceId()).toBeUndefined();
   });
 });

package/src/libs/observability/traceparent.ts

@@ -1,13 +1,5 @@
-import type {
-  Span,
-  Context as OtContext,
-  TextMapGetter
-} from '@lobechat/observability-otel/api';
-import {
-  context as otContext,
-  propagation,
-  trace,
-} from '@lobechat/observability-otel/api';
+import type { Context as OtContext, Span, TextMapGetter } from '@lobechat/observability-otel/api';
+import { context as otContext, propagation, trace } from '@lobechat/observability-otel/api';
 
 // NOTICE: do not try to optimize this into .repeat(...) or similar,
 // here served for better search / semantic search purpose for further diagnostic
@@ -47,6 +39,16 @@ export const getActiveTraceparent = () => {
   return toTraceparent(span as Span);
 };
 
+/**
+ * Get the traceId from the active span.
+ */
+export const getActiveTraceId = () => {
+  const span = trace.getActiveSpan();
+  if (!isValidContext(span)) return undefined;
+
+  return span!.spanContext().traceId;
+};
+
 /**
  * Injects the active context into headers using the configured propagator (W3C by default).
  * Also returns the traceparent for convenience.

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

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

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

@@ -14,14 +14,33 @@ import { desc, eq } from 'drizzle-orm';
 
 import { UserMemoryModel } from '@/database/models/userMemory';
 import { UserPersonaModel } from '@/database/models/userMemory/persona';
+import { AiInfraRepos } from '@/database/repositories/aiInfra';
 import { LobeChatDatabase } from '@/database/type';
 import {
   MemoryAgentConfig,
   parseMemoryExtractionConfig,
 } from '@/server/globalConfig/parseMemoryExtractionConfig';
+import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
 import { LayersEnum } from '@/types/userMemory';
 import { trimBasedOnBatchProbe } from '@/utils/chunkers';
 
+const extractCredentialsFromVault = (
+  vault?: Record<string, unknown>,
+): { apiKey?: string; baseURL?: string } => {
+  if (!vault || typeof vault !== 'object') return {};
+
+  const apiKey =
+    'apiKey' in vault && typeof (vault as any).apiKey === 'string'
+      ? (vault as any).apiKey
+      : undefined;
+  const baseURL =
+    'baseURL' in vault && typeof (vault as any).baseURL === 'string'
+      ? (vault as any).baseURL
+      : undefined;
+
+  return { apiKey, baseURL };
+};
+
 interface UserPersonaAgentPayload {
   existingPersona?: string | null;
   language?: string;
@@ -45,7 +64,6 @@ interface UserPersonaAgentResult {
 export class UserPersonaService {
   private readonly preferredLanguage?: string;
   private readonly db: LobeChatDatabase;
-  private readonly runtime: ModelRuntime;
   private readonly agentConfig: MemoryAgentConfig;
 
   constructor(db: LobeChatDatabase) {
@@ -54,13 +72,36 @@ export class UserPersonaService {
     this.db = db;
     this.preferredLanguage = agentPersonaWriter.language;
     this.agentConfig = agentPersonaWriter;
-    this.runtime = ModelRuntime.initializeWithProvider(agentPersonaWriter.provider || 'openai', {
-      apiKey: agentPersonaWriter.apiKey,
-      baseURL: agentPersonaWriter.baseURL,
-    });
   }
 
   async composeWriting(payload: UserPersonaAgentPayload): Promise<UserPersonaAgentResult> {
+    const aiInfraRepos = new AiInfraRepos(this.db, payload.userId, {});
+    const runtimeState = await aiInfraRepos.getAiProviderRuntimeState(
+      KeyVaultsGateKeeper.getUserKeyVaults,
+    );
+
+    const providerId = await AiInfraRepos.tryMatchingProviderFrom(runtimeState, {
+      fallbackProvider: this.agentConfig.provider,
+      label: 'persona writer',
+      modelId: this.agentConfig.model,
+    });
+
+    const normalizedProvider = providerId.toLowerCase();
+    const { apiKey: vaultApiKey, baseURL: vaultBaseURL } = extractCredentialsFromVault(
+      runtimeState.runtimeConfig?.[normalizedProvider]?.keyVaults,
+    );
+
+    const useVaultCredential = !!vaultApiKey;
+    const apiKey = useVaultCredential ? vaultApiKey : this.agentConfig.apiKey;
+    const baseURL = useVaultCredential
+      ? vaultBaseURL || this.agentConfig.baseURL
+      : this.agentConfig.baseURL;
+
+    const runtime = await ModelRuntime.initializeWithProvider(normalizedProvider, {
+      apiKey,
+      baseURL,
+    });
+
     const personaModel = new UserPersonaModel(this.db, payload.userId);
     const lastDocument = await personaModel.getLatestPersonaDocument();
     const existingPersonaBaseline = payload.existingPersona ?? lastDocument?.persona;
@@ -68,7 +109,7 @@ export class UserPersonaService {
     const extractor = new UserPersonaExtractor({
       agent: 'user-persona',
       model: this.agentConfig.model,
-      modelRuntime: this.runtime,
+      modelRuntime: runtime,
     });
 
     const agentResult = await extractor.toolCall({