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

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

@@ -4,6 +4,7 @@ import {
   DEFAULT_USER_MEMORY_EMBEDDING_DIMENSIONS,
   DEFAULT_USER_MEMORY_EMBEDDING_MODEL_ITEM,
 } from '@lobechat/const';
+import { type LobeChatDatabase } from '@lobechat/database';
 import {
   AddIdentityActionSchema,
   ContextMemoryItemSchema,
@@ -33,10 +34,9 @@ import {
   userMemoriesPreferences,
 } from '@/database/schemas';
 import { authedProcedure, router } from '@/libs/trpc/lambda';
-import { keyVaults, serverDatabase } from '@/libs/trpc/lambda/middleware';
+import { serverDatabase } from '@/libs/trpc/lambda/middleware';
 import { getServerDefaultFilesConfig } from '@/server/globalConfig';
-import { initModelRuntimeWithUserPayload } from '@/server/modules/ModelRuntime';
-import { type ClientSecretPayload } from '@/types/auth';
+import { initModelRuntimeFromDB } from '@/server/modules/ModelRuntime';
 
 const EMPTY_SEARCH_RESULT: SearchMemoryResult = {
   contexts: [],
@@ -45,8 +45,9 @@ const EMPTY_SEARCH_RESULT: SearchMemoryResult = {
 };
 
 type MemorySearchContext = {
-  jwtPayload: ClientSecretPayload;
   memoryModel: UserMemoryModel;
+  serverDB: LobeChatDatabase;
+  userId: string;
 };
 
 type MemorySearchResult = Awaited<ReturnType<UserMemoryModel['searchWithEmbedding']>>;
@@ -108,12 +109,12 @@ const searchUserMemories = async (
   ctx: MemorySearchContext,
   input: z.infer<typeof searchMemorySchema>,
 ): Promise<SearchMemoryResult> => {
-  const agentRuntime = await initModelRuntimeWithUserPayload(ModelProvider.OpenAI, ctx.jwtPayload);
-
-  const { model: embeddingModel } =
+  const { provider, model: embeddingModel } =
     getServerDefaultFilesConfig().embeddingModel || DEFAULT_USER_MEMORY_EMBEDDING_MODEL_ITEM;
+  // Read user's provider config from database
+  const modelRuntime = await initModelRuntimeFromDB(ctx.serverDB, ctx.userId, provider);
 
-  const queryEmbeddings = await agentRuntime.embeddings({
+  const queryEmbeddings = await modelRuntime.embeddings({
     dimensions: DEFAULT_USER_MEMORY_EMBEDDING_DIMENSIONS,
     input: input.query,
     model: embeddingModel,
@@ -133,11 +134,10 @@ const searchUserMemories = async (
   return mapMemorySearchResult(layeredResults);
 };
 
-const getEmbeddingRuntime = async (jwtPayload: ClientSecretPayload) => {
-  const agentRuntime = await initModelRuntimeWithUserPayload(
-    ENABLE_BUSINESS_FEATURES ? BRANDING_PROVIDER : ModelProvider.OpenAI,
-    jwtPayload,
-  );
+const getEmbeddingRuntime = async (serverDB: LobeChatDatabase, userId: string) => {
+  const provider = ENABLE_BUSINESS_FEATURES ? BRANDING_PROVIDER : ModelProvider.OpenAI;
+  // Read user's provider config from database
+  const agentRuntime = await initModelRuntimeFromDB(serverDB, userId, provider);
   const { model: embeddingModel } =
     getServerDefaultFilesConfig().embeddingModel || DEFAULT_USER_MEMORY_EMBEDDING_MODEL_ITEM;
 
@@ -197,17 +197,14 @@ const normalizeEmbeddable = (value?: string | null): string | undefined => {
   return trimmed.length > 0 ? trimmed : undefined;
 };
 
-const memoryProcedure = authedProcedure
-  .use(serverDatabase)
-  .use(keyVaults)
-  .use(async (opts) => {
-    const { ctx } = opts;
-    return opts.next({
-      ctx: {
-        memoryModel: new UserMemoryModel(ctx.serverDB, ctx.userId),
-      },
-    });
+const memoryProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
+  const { ctx } = opts;
+  return opts.next({
+    ctx: {
+      memoryModel: new UserMemoryModel(ctx.serverDB, ctx.userId),
+    },
   });
+});
 
 export const userMemoriesRouter = router({
   getMemoryDetail: memoryProcedure
@@ -316,7 +313,10 @@ export const userMemoriesRouter = router({
     .mutation(async ({ ctx, input }) => {
       try {
         const options = input ?? {};
-        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(ctx.jwtPayload);
+        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(
+          ctx.serverDB,
+          ctx.userId,
+        );
         const concurrency = options.concurrency ?? 10;
         const shouldProcess = (key: ReEmbedTableKey) =>
           !options.only || options.only.length === 0 || options.only.includes(key);
@@ -743,7 +743,10 @@ export const userMemoriesRouter = router({
     .input(ContextMemoryItemSchema)
     .mutation(async ({ input, ctx }) => {
       try {
-        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(ctx.jwtPayload);
+        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(
+          ctx.serverDB,
+          ctx.userId,
+        );
         const embed = createEmbedder(agentRuntime, embeddingModel);
 
         const summaryEmbedding = await embed(input.summary);
@@ -795,7 +798,10 @@ export const userMemoriesRouter = router({
     .input(ExperienceMemoryItemSchema)
     .mutation(async ({ input, ctx }) => {
       try {
-        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(ctx.jwtPayload);
+        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(
+          ctx.serverDB,
+          ctx.userId,
+        );
         const embed = createEmbedder(agentRuntime, embeddingModel);
 
         const summaryEmbedding = await embed(input.summary);
@@ -848,7 +854,10 @@ export const userMemoriesRouter = router({
     .input(AddIdentityActionSchema)
     .mutation(async ({ input, ctx }) => {
       try {
-        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(ctx.jwtPayload);
+        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(
+          ctx.serverDB,
+          ctx.userId,
+        );
         const embed = createEmbedder(agentRuntime, embeddingModel);
 
         const summaryEmbedding = await embed(input.summary);
@@ -913,7 +922,10 @@ export const userMemoriesRouter = router({
     .input(PreferenceMemoryItemSchema)
     .mutation(async ({ input, ctx }) => {
       try {
-        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(ctx.jwtPayload);
+        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(
+          ctx.serverDB,
+          ctx.userId,
+        );
         const embed = createEmbedder(agentRuntime, embeddingModel);
 
         const summaryEmbedding = await embed(input.summary);
@@ -1003,7 +1015,10 @@ export const userMemoriesRouter = router({
     .input(UpdateIdentityActionSchema)
     .mutation(async ({ input, ctx }) => {
       try {
-        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(ctx.jwtPayload);
+        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(
+          ctx.serverDB,
+          ctx.userId,
+        );
         const embed = createEmbedder(agentRuntime, embeddingModel);
 
         let summaryVector1024: number[] | null | undefined;

package/src/server/services/agentRuntime/AgentRuntimeService.test.ts

@@ -364,6 +364,93 @@ describe('AgentRuntimeService', () => {
       });
     });
 
+    it('should call onComplete with error in finalState when execution fails', async () => {
+      const error = new Error('Runtime error');
+      const mockRuntime = { step: vi.fn().mockRejectedValue(error) };
+      vi.spyOn(service as any, 'createAgentRuntime').mockReturnValue({ runtime: mockRuntime });
+
+      // Register onComplete callback
+      const mockOnComplete = vi.fn();
+      service.registerStepCallbacks('test-operation-1', {
+        onComplete: mockOnComplete,
+      });
+
+      await expect(service.executeStep(mockParams)).rejects.toThrow('Runtime error');
+
+      // Verify onComplete is called with error in finalState as ChatMessageError
+      // ChatErrorType.InternalServerError = 500
+      expect(mockOnComplete).toHaveBeenCalledWith(
+        expect.objectContaining({
+          operationId: 'test-operation-1',
+          reason: 'error',
+          finalState: expect.objectContaining({
+            error: expect.objectContaining({
+              type: 500, // ChatErrorType.InternalServerError
+              message: 'Runtime error',
+              body: expect.objectContaining({ name: 'Error' }),
+            }),
+          }),
+        }),
+      );
+    });
+
+    it('should call onComplete with ChatCompletionErrorPayload in finalState', async () => {
+      // Simulate LLM error format: { errorType: 'InvalidProviderAPIKey', error: { ... } }
+      const llmError = {
+        errorType: 'InvalidProviderAPIKey',
+        error: { status: 401 },
+        provider: 'openai',
+      };
+      const mockRuntime = { step: vi.fn().mockRejectedValue(llmError) };
+      vi.spyOn(service as any, 'createAgentRuntime').mockReturnValue({ runtime: mockRuntime });
+
+      // Register onComplete callback
+      const mockOnComplete = vi.fn();
+      service.registerStepCallbacks('test-operation-1', {
+        onComplete: mockOnComplete,
+      });
+
+      await expect(service.executeStep(mockParams)).rejects.toEqual(llmError);
+
+      // Verify error is formatted correctly with type from errorType
+      expect(mockOnComplete).toHaveBeenCalledWith(
+        expect.objectContaining({
+          operationId: 'test-operation-1',
+          reason: 'error',
+          finalState: expect.objectContaining({
+            error: expect.objectContaining({
+              type: 'InvalidProviderAPIKey',
+              message: 'InvalidProviderAPIKey',
+              body: expect.objectContaining({ status: 401 }),
+            }),
+          }),
+        }),
+      );
+    });
+
+    it('should save error state to coordinator for later retrieval (inMemory mode fix)', async () => {
+      const error = new Error('Test error for inMemory mode');
+      const mockRuntime = { step: vi.fn().mockRejectedValue(error) };
+      vi.spyOn(service as any, 'createAgentRuntime').mockReturnValue({ runtime: mockRuntime });
+
+      // Spy on coordinator.saveAgentState to verify it's called with error state
+      const saveStateSpy = vi.spyOn((service as any).coordinator, 'saveAgentState');
+
+      await expect(service.executeStep(mockParams)).rejects.toThrow('Test error for inMemory mode');
+
+      // Verify saveAgentState is called with error state before onComplete
+      expect(saveStateSpy).toHaveBeenCalledWith(
+        'test-operation-1',
+        expect.objectContaining({
+          error: expect.objectContaining({
+            type: 500, // ChatErrorType.InternalServerError
+            message: 'Test error for inMemory mode',
+          }),
+          status: 'error',
+        }),
+      );
+    });
+
     it('should handle human intervention', async () => {
       const paramsWithIntervention = {
         ...mockParams,

package/src/server/services/agentRuntime/AgentRuntimeService.ts

@@ -4,6 +4,7 @@ import {
   type AgentState,
   GeneralChatAgent,
 } from '@lobechat/agent-runtime';
+import { AgentRuntimeErrorType, ChatErrorType, type ChatMessageError } from '@lobechat/types';
 import debug from 'debug';
 import urlJoin from 'url-join';
 
@@ -41,6 +42,43 @@ import type {
 
 const log = debug('lobe-server:agent-runtime-service');
 
+/**
+ * Formats an error into ChatMessageError structure
+ * Handles various error formats from LLM execution and other sources
+ */
+function formatErrorForState(error: unknown): ChatMessageError {
+  // Handle ChatCompletionErrorPayload format from LLM errors
+  // e.g., { errorType: 'InvalidProviderAPIKey', error: { ... }, provider: 'openai' }
+  if (error && typeof error === 'object' && 'errorType' in error) {
+    const payload = error as {
+      error?: unknown;
+      errorType: ChatMessageError['type'];
+      message?: string;
+    };
+    return {
+      body: payload.error || error,
+      message: payload.message || String(payload.errorType),
+      type: payload.errorType,
+    };
+  }
+
+  // Handle standard Error objects
+  if (error instanceof Error) {
+    return {
+      body: { name: error.name },
+      message: error.message,
+      type: ChatErrorType.InternalServerError,
+    };
+  }
+
+  // Fallback for unknown error types
+  return {
+    body: error,
+    message: String(error),
+    type: AgentRuntimeErrorType.AgentRuntimeError,
+  };
+}
+
 export interface AgentRuntimeServiceOptions {
   /**
    * Coordinator configuration options
@@ -92,6 +130,7 @@ export class AgentRuntimeService {
 
     return urlJoin(baseUrl, '/api/agent');
   }
+  private serverDB: LobeChatDatabase;
   private userId: string;
   private messageModel: MessageModel;
 
@@ -107,6 +146,7 @@ export class AgentRuntimeService {
     });
     this.queueService =
       options?.queueService === null ? null : (options?.queueService ?? new QueueService());
+    this.serverDB = db;
     this.userId = userId;
     this.messageModel = new MessageModel(db, this.userId);
 
@@ -431,12 +471,22 @@ export class AgentRuntimeService {
         type: 'error',
       });
 
+      // Build and save error state so it's persisted for later retrieval
+      const errorState = await this.coordinator.loadAgentState(operationId);
+      const finalStateWithError = {
+        ...errorState!,
+        error: formatErrorForState(error),
+        status: 'error' as const,
+      };
+
+      // Save the error state to coordinator so getOperationStatus can retrieve it
+      await this.coordinator.saveAgentState(operationId, finalStateWithError);
+
       // Also call onComplete callback when execution fails
       if (callbacks?.onComplete) {
         try {
-          const errorState = await this.coordinator.loadAgentState(operationId);
           await callbacks.onComplete({
-            finalState: errorState!,
+            finalState: finalStateWithError,
             operationId,
             reason: 'error',
           });
@@ -794,6 +844,7 @@ export class AgentRuntimeService {
     const executorContext: RuntimeExecutorContext = {
       messageModel: this.messageModel,
       operationId,
+      serverDB: this.serverDB,
       stepIndex,
       streamManager: this.streamManager,
       toolExecutionService: this.toolExecutionService,

package/src/server/services/agentRuntime/__tests__/executeSync.test.ts

@@ -20,7 +20,7 @@ vi.mock('@/database/models/message', () => ({
 // Mock ModelRuntime
 vi.mock('@/server/modules/ModelRuntime', () => ({
   initializeRuntimeOptions: vi.fn(),
-  initModelRuntimeWithUserPayload: vi.fn().mockReturnValue({
+  initModelRuntimeFromDB: vi.fn().mockResolvedValue({
     chat: vi.fn(),
   }),
   ApiKeyManager: vi.fn().mockImplementation(() => ({
@@ -269,11 +269,7 @@ describe('AgentRuntimeService.executeSync', () => {
         await streamEventManager.publishAgentRuntimeEnd(operationId, 1, { status: 'done' });
       }, 10);
 
-      const event = await streamEventManager.waitForEvent(
-        operationId,
-        'agent_runtime_end',
-        1000,
-      );
+      const event = await streamEventManager.waitForEvent(operationId, 'agent_runtime_end', 1000);
 
       expect(event.type).toBe('agent_runtime_end');
     });

package/src/server/services/agentRuntime/__tests__/stepLifecycleCallbacks.test.ts

@@ -25,7 +25,7 @@ vi.mock('@/server/modules/ModelRuntime', () => ({
     getApiKey: vi.fn(),
   })),
   initializeRuntimeOptions: vi.fn(),
-  initModelRuntimeWithUserPayload: vi.fn().mockReturnValue({
+  initModelRuntimeFromDB: vi.fn().mockResolvedValue({
     chat: vi.fn(),
   }),
 }));

package/src/server/services/chunk/index.ts

@@ -1,5 +1,4 @@
 import { type LobeChatDatabase } from '@lobechat/database';
-import { type ClientSecretPayload } from '@lobechat/types';
 
 import { AsyncTaskModel } from '@/database/models/asyncTask';
 import { FileModel } from '@/database/models/file';
@@ -31,7 +30,7 @@ export class ChunkService {
     return this.chunkClient.chunkContent(params);
   }
 
-  async asyncEmbeddingFileChunks(fileId: string, payload: ClientSecretPayload) {
+  async asyncEmbeddingFileChunks(fileId: string) {
     const result = await this.fileModel.findById(fileId);
 
     if (!result) return;
@@ -44,7 +43,8 @@ export class ChunkService {
 
     await this.fileModel.update(fileId, { embeddingTaskId: asyncTaskId });
 
-    const asyncCaller = await createAsyncCaller({ jwtPayload: payload, userId: this.userId });
+    // Async router will read keyVaults from DB, no need to pass jwtPayload
+    const asyncCaller = await createAsyncCaller({ userId: this.userId });
 
     // trigger embedding task asynchronously
     try {
@@ -67,7 +67,7 @@ export class ChunkService {
   /**
    * parse file to chunks with async task
    */
-  async asyncParseFileToChunks(fileId: string, payload: ClientSecretPayload, skipExist?: boolean) {
+  async asyncParseFileToChunks(fileId: string, skipExist?: boolean) {
     const result = await this.fileModel.findById(fileId);
 
     if (!result) return;
@@ -83,7 +83,8 @@ export class ChunkService {
 
     await this.fileModel.update(fileId, { chunkTaskId: asyncTaskId });
 
-    const asyncCaller = await createAsyncCaller({ jwtPayload: payload, userId: this.userId });
+    // Async router will read keyVaults from DB, no need to pass jwtPayload
+    const asyncCaller = await createAsyncCaller({ userId: this.userId });
 
     // trigger parse file task asynchronously
     asyncCaller.file.parseFileToChunks({ fileId: fileId, taskId: asyncTaskId }).catch(async (e) => {

package/src/server/services/toolExecution/types.ts

@@ -1,10 +1,9 @@
 import { type LobeToolManifest } from '@lobechat/context-engine';
-import { type ChatToolPayload, type ClientSecretPayload } from '@lobechat/types';
+import { type ChatToolPayload } from '@lobechat/types';
 
 export interface ToolExecutionContext {
   toolManifestMap: Record<string, LobeToolManifest>;
   userId?: string;
-  userPayload?: ClientSecretPayload;
 }
 
 export interface ToolExecutionResult {

package/src/services/__tests__/_url.test.ts

@@ -17,7 +17,6 @@ describe('API_ENDPOINTS', () => {
     expect(API_ENDPOINTS.chat('openai')).toBe('/webapi/chat/openai');
     expect(API_ENDPOINTS.models('anthropic')).toBe('/webapi/models/anthropic');
     expect(API_ENDPOINTS.modelPull('azure')).toBe('/webapi/models/azure/pull');
-    expect(API_ENDPOINTS.images('dalle')).toBe('/webapi/text-to-image/dalle');
     expect(API_ENDPOINTS.tts('openai')).toBe('/webapi/tts/openai');
   });
 });

package/src/services/_url.ts

@@ -21,9 +21,6 @@ export const API_ENDPOINTS = {
   modelPull: (provider: string) =>
     withElectronProtocolIfElectron(`/webapi/models/${provider}/pull`),
 
-  // image
-  images: (provider: string) => withElectronProtocolIfElectron(`/webapi/text-to-image/${provider}`),
-
   // STT
   stt: withElectronProtocolIfElectron('/webapi/stt/openai'),
 

package/src/services/aiChat.ts

@@ -2,7 +2,6 @@ import { type SendMessageServerParams, type StructureOutputParams } from '@lobec
 import { cleanObject } from '@lobechat/utils';
 
 import { lambdaClient } from '@/libs/trpc/client';
-import { createXorKeyVaultsPayload } from '@/services/_auth';
 
 class AiChatService {
   sendMessageInServer = async (
@@ -15,17 +14,11 @@ class AiChatService {
     });
   };
 
-  generateJSON = async (
-    params: Omit<StructureOutputParams, 'keyVaultsPayload'>,
-    abortController: AbortController,
-  ) => {
-    return lambdaClient.aiChat.outputJSON.mutate(
-      { ...params, keyVaultsPayload: createXorKeyVaultsPayload(params.provider) },
-      {
-        context: { showNotification: false },
-        signal: abortController?.signal,
-      },
-    );
+  generateJSON = async (params: StructureOutputParams, abortController: AbortController) => {
+    return lambdaClient.aiChat.outputJSON.mutate(params, {
+      context: { showNotification: false },
+      signal: abortController?.signal,
+    });
   };
 
   // sendGroupMessageInServer = async (params: SendMessageServerParams) => {

package/src/store/chat/slices/aiChat/actions/streamingExecutor.ts

@@ -164,8 +164,6 @@ export const streamingExecutor: StateCreator<
     const operation = operationId ? get().operations[operationId] : undefined;
     const scope = operation?.context.scope;
 
-    console.log('[internal_createAgentState] Operation scope:', { operationId, scope });
-
     // Resolve agent config with builtin agent runtime config merged
     // This ensures runtime plugins (e.g., 'lobe-agent-builder' for Agent Builder) are included
     const { agentConfig: agentConfigData, plugins: pluginIds } = resolveAgentConfig({

package/src/app/(backend)/webapi/text-to-image/[provider]/route.ts

@@ -1,74 +0,0 @@
-import { type ChatCompletionErrorPayload, type TextToImagePayload } from '@lobechat/model-runtime';
-import { ChatErrorType } from '@lobechat/types';
-import { NextResponse } from 'next/server';
-
-import { checkAuth } from '@/app/(backend)/middleware/auth';
-import { initModelRuntimeWithUserPayload } from '@/server/modules/ModelRuntime';
-import { createErrorResponse } from '@/utils/errorResponse';
-
-export const preferredRegion = [
-  'arn1',
-  'bom1',
-  'cdg1',
-  'cle1',
-  'cpt1',
-  'dub1',
-  'fra1',
-  'gru1',
-  'hnd1',
-  'iad1',
-  'icn1',
-  'kix1',
-  'lhr1',
-  'pdx1',
-  'sfo1',
-  'sin1',
-  'syd1',
-];
-
-// return NextResponse.json(
-//   {
-//     body: {
-//       endpoint: 'https://ai****ix.com/v1',
-//       error: {
-//         code: 'content_policy_violation',
-//         message:
-//           'Your request was rejected as a result of our safety system. Image descriptions generated from your prompt may contain text that is not allowed by our safety system. If you believe this was done in error, your request may succeed if retried, or by adjusting your prompt.',
-//         param: null,
-//         type: 'invalid_request_error',
-//       },
-//       provider: 'openai',
-//     },
-//     errorType: 'OpenAIBizError',
-//   },
-//   { status: 400 },
-// );
-
-export const POST = checkAuth(async (req: Request, { params, jwtPayload }) => {
-  const provider = (await params)!.provider!;
-
-  try {
-    // ============  1. init chat model   ============ //
-    const agentRuntime = await initModelRuntimeWithUserPayload(provider, jwtPayload);
-
-    // ============  2. create chat completion   ============ //
-
-    const data = (await req.json()) as TextToImagePayload;
-
-    const images = await agentRuntime.textToImage(data);
-
-    return NextResponse.json(images);
-  } catch (e) {
-    const {
-      errorType = ChatErrorType.InternalServerError,
-      error: errorContent,
-      ...res
-    } = e as ChatCompletionErrorPayload;
-
-    const error = errorContent || e;
-    // track the error at server side
-    console.error(`Route: [${provider}] ${errorType}:`, error);
-
-    return createErrorResponse(errorType, { error, ...res, provider });
-  }
-});