@lobehub/lobehub 2.0.0-next.201 → 2.0.0-next.203

package/src/server/modules/ModelRuntime/index.ts

@@ -1,16 +1,113 @@
 import { type GoogleGenAIOptions } from '@google/genai';
 import { ModelRuntime } from '@lobechat/model-runtime';
 import { LobeVertexAI } from '@lobechat/model-runtime/vertexai';
-import { type ClientSecretPayload } from '@lobechat/types';
+import {
+  type AWSBedrockKeyVault,
+  type AzureOpenAIKeyVault,
+  type ClientSecretPayload,
+  type CloudflareKeyVault,
+  type ComfyUIKeyVault,
+  type OpenAICompatibleKeyVault,
+  type VertexAIKeyVault,
+} from '@lobechat/types';
 import { safeParseJSON } from '@lobechat/utils';
 import { ModelProvider } from 'model-bank';
 
+import { AiProviderModel } from '@/database/models/aiProvider';
+import { type LobeChatDatabase } from '@/database/type';
 import { getLLMConfig } from '@/envs/llm';
 
+import { KeyVaultsGateKeeper } from '../KeyVaultsEncrypt';
 import apiKeyManager from './apiKeyManager';
 
 export * from './trace';
 
+/**
+ * Combined KeyVaults type for all providers
+ */
+type ProviderKeyVaults = OpenAICompatibleKeyVault &
+  AzureOpenAIKeyVault &
+  AWSBedrockKeyVault &
+  CloudflareKeyVault &
+  ComfyUIKeyVault &
+  VertexAIKeyVault;
+
+/**
+ * Build ClientSecretPayload from keyVaults stored in database
+ *
+ * This is the server-side equivalent of the frontend's getProviderAuthPayload function.
+ * It converts the keyVaults object from database to the ClientSecretPayload format
+ * expected by initModelRuntimeWithUserPayload.
+ *
+ * @param provider - The model provider
+ * @param keyVaults - The keyVaults object from database (already decrypted)
+ * @returns ClientSecretPayload for the provider
+ */
+export const buildPayloadFromKeyVaults = (
+  provider: string,
+  keyVaults: ProviderKeyVaults,
+): ClientSecretPayload => {
+  switch (provider) {
+    case ModelProvider.Bedrock: {
+      const { accessKeyId, region, secretAccessKey, sessionToken } = keyVaults;
+      const apiKey = (secretAccessKey || '') + (accessKeyId || '');
+
+      return {
+        apiKey,
+        awsAccessKeyId: accessKeyId,
+        awsRegion: region,
+        awsSecretAccessKey: secretAccessKey,
+        awsSessionToken: sessionToken,
+      };
+    }
+
+    case ModelProvider.Azure: {
+      return {
+        apiKey: keyVaults.apiKey,
+        azureApiVersion: keyVaults.apiVersion,
+        baseURL: keyVaults.baseURL || keyVaults.endpoint,
+      };
+    }
+
+    case ModelProvider.Ollama: {
+      return { baseURL: keyVaults.baseURL };
+    }
+
+    case ModelProvider.Cloudflare: {
+      return {
+        apiKey: keyVaults.apiKey,
+        cloudflareBaseURLOrAccountID: keyVaults.baseURLOrAccountID,
+      };
+    }
+
+    case ModelProvider.ComfyUI: {
+      return {
+        apiKey: keyVaults.apiKey,
+        authType: keyVaults.authType,
+        baseURL: keyVaults.baseURL,
+        customHeaders: keyVaults.customHeaders,
+        password: keyVaults.password,
+        username: keyVaults.username,
+      };
+    }
+
+    case ModelProvider.VertexAI: {
+      return {
+        apiKey: keyVaults.apiKey,
+        baseURL: keyVaults.baseURL,
+        vertexAIRegion: keyVaults.region,
+      };
+    }
+
+    default: {
+      return {
+        apiKey: keyVaults.apiKey,
+        baseURL: keyVaults.baseURL,
+      };
+    }
+  }
+};
+
 /**
  * Retrieves the options object from environment and apikeymanager
  * based on the provider and payload.
@@ -220,3 +317,43 @@ export const initModelRuntimeWithUserPayload = (
     ...params,
   });
 };
+
+/**
+ * Initialize ModelRuntime by reading user's provider configuration from database
+ *
+ * This function replaces the pattern of passing userPayload from frontend.
+ * It reads the user's AI provider configuration from the database, decrypts
+ * the keyVaults, and initializes the ModelRuntime.
+ *
+ * @param db - The database instance
+ * @param userId - The user ID
+ * @param provider - The model provider (e.g., 'openai', 'azure')
+ * @returns Promise<ModelRuntime> - The initialized ModelRuntime instance
+ *
+ * @example
+ * ```typescript
+ * const modelRuntime = await initModelRuntimeFromDB(db, userId, 'openai');
+ * const response = await modelRuntime.chat({ messages, model });
+ * ```
+ */
+export const initModelRuntimeFromDB = async (
+  db: LobeChatDatabase,
+  userId: string,
+  provider: string,
+): Promise<ModelRuntime> => {
+  // 1. Get user's provider configuration from database
+  const aiProviderModel = new AiProviderModel(db, userId);
+
+  // Use getAiProviderById with KeyVaultsGateKeeper.getUserKeyVaults as decryptor
+  const providerConfig = await aiProviderModel.getAiProviderById(
+    provider,
+    KeyVaultsGateKeeper.getUserKeyVaults,
+  );
+
+  // 2. Build ClientSecretPayload from keyVaults
+  const keyVaults = (providerConfig?.keyVaults || {}) as ProviderKeyVaults;
+  const payload = buildPayloadFromKeyVaults(provider, keyVaults);
+
+  // 3. Initialize ModelRuntime with the payload
+  return initModelRuntimeWithUserPayload(provider, payload);
+};

package/src/server/routers/async/__tests__/caller.test.ts

@@ -65,7 +65,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
       mockAppEnv.APP_URL = 'https://public.example.com';
       mockAppEnv.INTERNAL_APP_URL = 'http://localhost:3210';
 
-      await createAsyncServerClient('user-123', { apiKey: 'test-key' });
+      await createAsyncServerClient('user-123');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -80,7 +80,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
       mockAppEnv.APP_URL = 'https://fallback.example.com';
       mockAppEnv.INTERNAL_APP_URL = 'https://fallback.example.com'; // getInternalAppUrl() returns APP_URL
 
-      await createAsyncServerClient('user-456', {});
+      await createAsyncServerClient('user-456');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -92,7 +92,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
       mockAppEnv.APP_URL = 'https://cdn-proxied.example.com';
       mockAppEnv.INTERNAL_APP_URL = 'http://127.0.0.1:3210';
 
-      await createAsyncServerClient('user-789', {});
+      await createAsyncServerClient('user-789');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -105,7 +105,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
       mockAppEnv.APP_URL = 'https://public.example.com';
       mockAppEnv.INTERNAL_APP_URL = 'http://lobe-service:3210';
 
-      await createAsyncServerClient('user-docker', {});
+      await createAsyncServerClient('user-docker');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -116,7 +116,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
     it('should handle INTERNAL_APP_URL with trailing slash', async () => {
       mockAppEnv.INTERNAL_APP_URL = 'http://localhost:3210/';
 
-      await createAsyncServerClient('user-trailing', {});
+      await createAsyncServerClient('user-trailing');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -128,7 +128,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
     it('should handle INTERNAL_APP_URL without trailing slash', async () => {
       mockAppEnv.INTERNAL_APP_URL = 'https://example.com';
 
-      await createAsyncServerClient('user-no-trailing', {});
+      await createAsyncServerClient('user-no-trailing');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -139,7 +139,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
 
   describe('authentication and headers', () => {
     it('should include Authorization header with internal JWT token', async () => {
-      await createAsyncServerClient('user-auth', {});
+      await createAsyncServerClient('user-auth');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -148,19 +148,16 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
       expect(httpLinkOptions.headers.Authorization).toBe('mock-internal-jwt-token');
     });
 
-    it('should encrypt and include user payload in x-lobe-chat-auth header', async () => {
-      const testPayload = { apiKey: 'test-api-key-value', provider: 'openai' };
+    it('should encrypt and include userId in x-lobe-chat-auth header', async () => {
       const mockEncrypt = vi.fn().mockResolvedValue('test-encrypted-auth-data');
       vi.mocked(KeyVaultsGateKeeper.initWithEnvKey).mockResolvedValueOnce({
         encrypt: mockEncrypt,
       } as any);
 
-      await createAsyncServerClient('user-encrypt', testPayload);
+      await createAsyncServerClient('user-encrypt');
 
       expect(KeyVaultsGateKeeper.initWithEnvKey).toHaveBeenCalled();
-      expect(mockEncrypt).toHaveBeenCalledWith(
-        JSON.stringify({ payload: testPayload, userId: 'user-encrypt' }),
-      );
+      expect(mockEncrypt).toHaveBeenCalledWith(JSON.stringify({ userId: 'user-encrypt' }));
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -176,7 +173,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
       const originalEnv = process.env.VERCEL_AUTOMATION_BYPASS_SECRET;
       process.env.VERCEL_AUTOMATION_BYPASS_SECRET = 'test-bypass-value';
 
-      await createAsyncServerClient('user-vercel', {});
+      await createAsyncServerClient('user-vercel');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -195,7 +192,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
     it('should not include Vercel bypass secret when not available', async () => {
       delete process.env.VERCEL_AUTOMATION_BYPASS_SECRET;
 
-      await createAsyncServerClient('user-no-vercel', {});
+      await createAsyncServerClient('user-no-vercel');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -211,9 +208,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
         encrypt: mockEncrypt,
       } as any);
 
-      await expect(createAsyncServerClient('user-enc-fail', {})).rejects.toThrow(
-        'Encryption failed',
-      );
+      await expect(createAsyncServerClient('user-enc-fail')).rejects.toThrow('Encryption failed');
 
       expect(KeyVaultsGateKeeper.initWithEnvKey).toHaveBeenCalled();
     });
@@ -223,7 +218,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
       mockAppEnv.APP_URL = 'https://only-app-url.com';
       mockAppEnv.INTERNAL_APP_URL = 'https://only-app-url.com'; // Result of fallback
 
-      await createAsyncServerClient('user-null', {});
+      await createAsyncServerClient('user-null');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -236,7 +231,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
       mockAppEnv.APP_URL = 'https://fallback-from-empty.com';
       mockAppEnv.INTERNAL_APP_URL = '';
 
-      await createAsyncServerClient('user-empty', {});
+      await createAsyncServerClient('user-empty');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -249,7 +244,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
     it('should handle malformed URL gracefully', async () => {
       mockAppEnv.INTERNAL_APP_URL = 'not-a-valid-url';
 
-      await createAsyncServerClient('user-malformed', {});
+      await createAsyncServerClient('user-malformed');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -262,7 +257,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
 
   describe('TRPC client configuration', () => {
     it('should configure httpLink with proper options', async () => {
-      await createAsyncServerClient('user-config', {});
+      await createAsyncServerClient('user-config');
 
       expect(httpLink).toHaveBeenCalled();
       const httpLinkOptions = vi.mocked(httpLink).mock.calls[0][0];
@@ -273,7 +268,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
     });
 
     it('should pass httpLink result to createTRPCClient', async () => {
-      await createAsyncServerClient('user-link', {});
+      await createAsyncServerClient('user-link');
 
       expect(createTRPCClient).toHaveBeenCalledWith({
         links: expect.arrayContaining([
@@ -286,7 +281,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
     });
 
     it('should return the created TRPC client', async () => {
-      const client = await createAsyncServerClient('user-return', {});
+      const client = await createAsyncServerClient('user-return');
 
       expect(client).toBeDefined();
       expect(client).toHaveProperty('_mockClient');
@@ -299,7 +294,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
       mockAppEnv.APP_URL = 'https://lobechat.example.com';
       mockAppEnv.INTERNAL_APP_URL = 'http://localhost:3210';
 
-      await createAsyncServerClient('prod-user', { apiKey: 'test-key' });
+      await createAsyncServerClient('prod-user');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -312,7 +307,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
       mockAppEnv.APP_URL = 'https://public.example.com';
       mockAppEnv.INTERNAL_APP_URL = 'http://lobe-chat-database:3210';
 
-      await createAsyncServerClient('docker-user', {});
+      await createAsyncServerClient('docker-user');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;
@@ -325,7 +320,7 @@ describe('createAsyncServerClient - INTERNAL_APP_URL Tests', () => {
       mockAppEnv.APP_URL = 'https://direct-access.example.com';
       mockAppEnv.INTERNAL_APP_URL = 'https://direct-access.example.com'; // Result of getInternalAppUrl() fallback
 
-      await createAsyncServerClient('direct-user', {});
+      await createAsyncServerClient('direct-user');
 
       const config = vi.mocked(createTRPCClient).mock.calls[0][0];
       const httpLinkOptions = config.links[0] as any;

package/src/server/routers/async/caller.ts

@@ -1,4 +1,3 @@
-import { type ClientSecretPayload } from '@lobechat/types';
 import { createTRPCClient, httpLink } from '@trpc/client';
 import superjson from 'superjson';
 import urlJoin from 'url-join';
@@ -12,12 +11,12 @@ import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
 import { asyncRouter } from './index';
 import type { AsyncRouter } from './index';
 
-export const createAsyncServerClient = async (userId: string, payload: ClientSecretPayload) => {
+export const createAsyncServerClient = async (userId: string) => {
   const token = await signInternalJWT();
   const gateKeeper = await KeyVaultsGateKeeper.initWithEnvKey();
   const headers: Record<string, string> = {
     Authorization: token,
-    [LOBE_CHAT_AUTH_HEADER]: await gateKeeper.encrypt(JSON.stringify({ payload, userId })),
+    [LOBE_CHAT_AUTH_HEADER]: await gateKeeper.encrypt(JSON.stringify({ userId })),
   };
 
   if (process.env.VERCEL_AUTOMATION_BYPASS_SECRET) {
@@ -49,7 +48,6 @@ const helperFunc = () => {
 export type UnifiedAsyncCaller = ReturnType<typeof helperFunc>;
 
 interface CreateCallerOptions {
-  jwtPayload: any;
   userId: string;
 }
 
@@ -60,9 +58,9 @@ interface CreateCallerOptions {
 export const createAsyncCaller = async (
   options: CreateCallerOptions,
 ): Promise<UnifiedAsyncCaller> => {
-  const { userId, jwtPayload } = options;
+  const { userId } = options;
 
-  const httpClient = await createAsyncServerClient(userId, jwtPayload);
+  const httpClient = await createAsyncServerClient(userId);
   const createRecursiveProxy = (client: any, path: string[]): any => {
     // The target is a dummy function, so that 'apply' can be triggered.
     return new Proxy(() => {}, {

package/src/server/routers/async/file.ts

@@ -15,7 +15,7 @@ import { type NewChunkItem, type NewEmbeddingsItem } from '@/database/schemas';
 import { fileEnv } from '@/envs/file';
 import { asyncAuthedProcedure, asyncRouter as router } from '@/libs/trpc/async';
 import { getServerDefaultFilesConfig } from '@/server/globalConfig';
-import { initModelRuntimeWithUserPayload } from '@/server/modules/ModelRuntime';
+import { initModelRuntimeFromDB } from '@/server/modules/ModelRuntime';
 import { ChunkService } from '@/server/services/chunk';
 import { FileService } from '@/server/services/file';
 import {
@@ -95,9 +95,14 @@ export const fileRouter = router({
             await pMap(
               requestArray,
               async (chunks) => {
-                const agentRuntime = initModelRuntimeWithUserPayload(provider, ctx.jwtPayload);
-
-                const embeddings = await agentRuntime.embeddings({
+                // Read user's provider config from database
+                const modelRuntime = await initModelRuntimeFromDB(
+                  ctx.serverDB,
+                  ctx.userId,
+                  provider,
+                );
+
+                const embeddings = await modelRuntime.embeddings({
                   dimensions: 1024,
                   input: chunks.map((c) => c.text),
                   model,
@@ -243,7 +248,7 @@ export const fileRouter = router({
 
           // if enable auto embedding, trigger the embedding task
           if (fileEnv.CHUNKS_AUTO_EMBEDDING) {
-            await chunkService.asyncEmbeddingFileChunks(input.fileId, ctx.jwtPayload);
+            await chunkService.asyncEmbeddingFileChunks(input.fileId);
           }
 
           return { success: true };

package/src/server/routers/async/image.ts

@@ -14,7 +14,7 @@ import { FileModel } from '@/database/models/file';
 import { GenerationModel } from '@/database/models/generation';
 import { GenerationBatchModel } from '@/database/models/generationBatch';
 import { asyncAuthedProcedure, asyncRouter as router } from '@/libs/trpc/async';
-import { initModelRuntimeWithUserPayload } from '@/server/modules/ModelRuntime';
+import { initModelRuntimeFromDB } from '@/server/modules/ModelRuntime';
 import { GenerationService } from '@/server/services/generation';
 
 const log = debug('lobe-image:async');
@@ -241,12 +241,13 @@ export const imageRouter = router({
         const imageGenerationPromise = async (signal: AbortSignal) => {
           log('Initializing agent runtime for provider: %s', provider);
 
-          const agentRuntime = initModelRuntimeWithUserPayload(provider, ctx.jwtPayload);
+          // Read user's provider config from database
+          const modelRuntime = await initModelRuntimeFromDB(ctx.serverDB, ctx.userId, provider);
 
           // Check if operation has been cancelled
           checkAbortSignal(signal);
           log('Agent runtime initialized, calling createImage');
-          const response = await agentRuntime.createImage!({
+          const response = await modelRuntime.createImage!({
             model,
             params: params as unknown as RuntimeImageGenParams,
           });
@@ -298,7 +299,7 @@ export const imageRouter = router({
           if (provider === 'comfyui') {
             // Use the public interface method to get auth headers
             // This avoids accessing private members and exposing credentials
-            authHeaders = agentRuntime.getAuthHeaders();
+            authHeaders = modelRuntime.getAuthHeaders();
             if (authHeaders) {
               log('Using authentication headers for ComfyUI image download');
             } else {

package/src/server/routers/async/ragEval.ts

@@ -15,7 +15,7 @@ import {
   EvaluationRecordModel,
 } from '@/database/server/models/ragEval';
 import { asyncAuthedProcedure, asyncRouter as router } from '@/libs/trpc/async';
-import { initModelRuntimeWithUserPayload } from '@/server/modules/ModelRuntime';
+import { initModelRuntimeFromDB } from '@/server/modules/ModelRuntime';
 import { ChunkService } from '@/server/services/chunk';
 import { AsyncTaskError } from '@/types/asyncTask';
 
@@ -51,9 +51,11 @@ export const ragEvalRouter = router({
 
       const now = Date.now();
       try {
-        const agentRuntime = await initModelRuntimeWithUserPayload(
+        // Read user's provider config from database
+        const modelRuntime = await initModelRuntimeFromDB(
+          ctx.serverDB,
+          ctx.userId,
           ModelProvider.OpenAI,
-          ctx.jwtPayload,
         );
 
         const { question, languageModel, embeddingModel } = evalRecord;
@@ -63,7 +65,7 @@ export const ragEvalRouter = router({
 
         // 如果不存在 questionEmbeddingId，那么就需要做一次 embedding
         if (!questionEmbeddingId) {
-          const embeddings = await agentRuntime.embeddings({
+          const embeddings = await modelRuntime.embeddings({
             dimensions: 1024,
             input: question,
             model: !!embeddingModel ? embeddingModel : DEFAULT_EMBEDDING_MODEL,
@@ -100,7 +102,7 @@ export const ragEvalRouter = router({
         // 做一次生成 LLM 答案生成
         const { messages } = chainAnswerWithContext({ context, knowledge: [], question });
 
-        const response = await agentRuntime.chat({
+        const response = await modelRuntime.chat({
           messages: messages!,
           model: !!languageModel ? languageModel : DEFAULT_MODEL,
           responseMode: 'json',

package/src/server/routers/lambda/__tests__/aiChat.test.ts

@@ -18,11 +18,8 @@ vi.mock('@/server/services/aiChat');
 vi.mock('@/server/services/file', () => ({
   FileService: vi.fn(),
 }));
-vi.mock('@/utils/server', () => ({
-  getXorPayload: vi.fn(),
-}));
 vi.mock('@/server/modules/ModelRuntime', () => ({
-  initModelRuntimeWithUserPayload: vi.fn(),
+  initModelRuntimeFromDB: vi.fn(),
 }));
 
 describe('aiChatRouter', () => {
@@ -690,22 +687,18 @@ describe('aiChatRouter', () => {
 
   describe('outputJSON', () => {
     it('should successfully generate structured output', async () => {
-      const { getXorPayload } = await import('@/utils/server');
-      const { initModelRuntimeWithUserPayload } = await import('@/server/modules/ModelRuntime');
+      const { initModelRuntimeFromDB } = await import('@/server/modules/ModelRuntime');
 
-      const mockPayload = { apiKey: 'test-key' };
       const mockResult = { object: { name: 'John', age: 30 } };
       const mockGenerateObject = vi.fn().mockResolvedValue(mockResult);
 
-      vi.mocked(getXorPayload).mockReturnValue(mockPayload);
-      vi.mocked(initModelRuntimeWithUserPayload).mockReturnValue({
+      vi.mocked(initModelRuntimeFromDB).mockResolvedValue({
         generateObject: mockGenerateObject,
       } as any);
 
-      const caller = aiChatRouter.createCaller(mockCtx as any);
+      const caller = aiChatRouter.createCaller({ ...mockCtx, serverDB: {} } as any);
 
       const input = {
-        keyVaultsPayload: 'encrypted-payload',
         messages: [{ content: 'test', role: 'user' }],
         model: 'gpt-4o',
         provider: 'openai',
@@ -720,8 +713,7 @@ describe('aiChatRouter', () => {
 
       const result = await caller.outputJSON(input);
 
-      expect(getXorPayload).toHaveBeenCalledWith('encrypted-payload');
-      expect(initModelRuntimeWithUserPayload).toHaveBeenCalledWith('openai', mockPayload);
+      expect(initModelRuntimeFromDB).toHaveBeenCalledWith({}, 'u1', 'openai');
       expect(mockGenerateObject).toHaveBeenCalledWith({
         messages: input.messages,
         model: 'gpt-4o',
@@ -731,28 +723,9 @@ describe('aiChatRouter', () => {
       expect(result).toEqual(mockResult);
     });
 
-    it('should throw error when keyVaultsPayload is invalid', async () => {
-      const { getXorPayload } = await import('@/utils/server');
-
-      vi.mocked(getXorPayload).mockReturnValue(undefined as any);
-
-      const caller = aiChatRouter.createCaller(mockCtx as any);
-
-      const input = {
-        keyVaultsPayload: 'invalid-payload',
-        messages: [],
-        model: 'gpt-4o',
-        provider: 'openai',
-      };
-
-      await expect(caller.outputJSON(input)).rejects.toThrow('keyVaultsPayload is not correct');
-    });
-
     it('should handle tools parameter when provided', async () => {
-      const { getXorPayload } = await import('@/utils/server');
-      const { initModelRuntimeWithUserPayload } = await import('@/server/modules/ModelRuntime');
+      const { initModelRuntimeFromDB } = await import('@/server/modules/ModelRuntime');
 
-      const mockPayload = { apiKey: 'test-key' };
       const mockTools = [
         {
           type: 'function' as const,
@@ -767,15 +740,13 @@ describe('aiChatRouter', () => {
       ];
       const mockGenerateObject = vi.fn().mockResolvedValue({ object: {} });
 
-      vi.mocked(getXorPayload).mockReturnValue(mockPayload);
-      vi.mocked(initModelRuntimeWithUserPayload).mockReturnValue({
+      vi.mocked(initModelRuntimeFromDB).mockResolvedValue({
         generateObject: mockGenerateObject,
       } as any);
 
-      const caller = aiChatRouter.createCaller(mockCtx as any);
+      const caller = aiChatRouter.createCaller({ ...mockCtx, serverDB: {} } as any);
 
       const input = {
-        keyVaultsPayload: 'encrypted-payload',
         messages: [],
         model: 'gpt-4o',
         provider: 'openai',

package/src/server/routers/lambda/aiChat.ts

@@ -3,7 +3,6 @@ import {
   type SendMessageServerResponse,
   StructureOutputSchema,
 } from '@lobechat/types';
-import { TRPCError } from '@trpc/server';
 import debug from 'debug';
 
 import { LOADING_FLAT } from '@/const/message';
@@ -13,11 +12,10 @@ import { ThreadModel } from '@/database/models/thread';
 import { TopicModel } from '@/database/models/topic';
 import { authedProcedure, router } from '@/libs/trpc/lambda';
 import { serverDatabase } from '@/libs/trpc/lambda/middleware';
-import { initModelRuntimeWithUserPayload } from '@/server/modules/ModelRuntime';
+import { initModelRuntimeFromDB } from '@/server/modules/ModelRuntime';
 import { resolveContext } from '@/server/routers/lambda/_helpers/resolveContext';
 import { AiChatService } from '@/server/services/aiChat';
 import { FileService } from '@/server/services/file';
-import { getXorPayload } from '@/utils/server';
 
 const log = debug('lobe-lambda-router:ai-chat');
 
@@ -37,28 +35,14 @@ const aiChatProcedure = authedProcedure.use(serverDatabase).use(async (opts) =>
 });
 
 export const aiChatRouter = router({
-  outputJSON: aiChatProcedure.input(StructureOutputSchema).mutation(async ({ input }) => {
+  outputJSON: aiChatProcedure.input(StructureOutputSchema).mutation(async ({ input, ctx }) => {
     log('outputJSON called with provider: %s, model: %s', input.provider, input.model);
     log('messages count: %d', input.messages.length);
     log('schema: %O', input.schema);
 
-    let payload: object | undefined;
-
-    try {
-      payload = getXorPayload(input.keyVaultsPayload);
-      log('payload parsed successfully');
-    } catch (e) {
-      log('payload parse error: %O', e);
-      console.warn('user payload parse error', e);
-    }
-
-    if (!payload) {
-      log('payload is empty, throwing error');
-      throw new TRPCError({ code: 'BAD_REQUEST', message: 'keyVaultsPayload is not correct' });
-    }
-
-    log('initializing model runtime with provider: %s', input.provider);
-    const modelRuntime = initModelRuntimeWithUserPayload(input.provider, payload);
+    log('initializing model runtime from DB with provider: %s', input.provider);
+    // Read user's provider config from database
+    const modelRuntime = await initModelRuntimeFromDB(ctx.serverDB, ctx.userId, input.provider);
 
     log('calling generateObject');
     const result = await modelRuntime.generateObject({