@lobehub/lobehub 2.0.0-next.4 → 2.0.0-next.40

package/packages/database/src/models/agent.ts

@@ -28,19 +28,21 @@ export class AgentModel {
   };
 
   getAgentAssignedKnowledge = async (id: string) => {
-    const knowledgeBaseResult = await this.db
-      .select({ enabled: agentsKnowledgeBases.enabled, knowledgeBases })
-      .from(agentsKnowledgeBases)
-      .where(eq(agentsKnowledgeBases.agentId, id))
-      .orderBy(desc(agentsKnowledgeBases.createdAt))
-      .leftJoin(knowledgeBases, eq(knowledgeBases.id, agentsKnowledgeBases.knowledgeBaseId));
-
-    const fileResult = await this.db
-      .select({ enabled: agentsFiles.enabled, files })
-      .from(agentsFiles)
-      .where(eq(agentsFiles.agentId, id))
-      .orderBy(desc(agentsFiles.createdAt))
-      .leftJoin(files, eq(files.id, agentsFiles.fileId));
+    // Run both queries in parallel for better performance
+    const [knowledgeBaseResult, fileResult] = await Promise.all([
+      this.db
+        .select({ enabled: agentsKnowledgeBases.enabled, knowledgeBases })
+        .from(agentsKnowledgeBases)
+        .where(eq(agentsKnowledgeBases.agentId, id))
+        .orderBy(desc(agentsKnowledgeBases.createdAt))
+        .leftJoin(knowledgeBases, eq(knowledgeBases.id, agentsKnowledgeBases.knowledgeBaseId)),
+      this.db
+        .select({ enabled: agentsFiles.enabled, files })
+        .from(agentsFiles)
+        .where(eq(agentsFiles.agentId, id))
+        .orderBy(desc(agentsFiles.createdAt))
+        .leftJoin(files, eq(files.id, agentsFiles.fileId)),
+    ]);
 
     return {
       files: fileResult.map((item) => ({
@@ -61,6 +63,7 @@ export class AgentModel {
     const item = await this.db.query.agentsToSessions.findFirst({
       where: eq(agentsToSessions.sessionId, sessionId),
     });
+
     if (!item) return;
 
     const agentId = item.agentId;

package/packages/database/src/models/message.ts

@@ -1,3 +1,4 @@
+import { INBOX_SESSION_ID } from '@lobechat/const';
 import {
   ChatFileItem,
   ChatImageItem,
@@ -6,7 +7,6 @@ import {
   ChatTranslate,
   ChatVideoItem,
   CreateMessageParams,
-  CreateMessageResult,
   DBMessageItem,
   ModelRankItem,
   NewMessageQueryParams,
@@ -14,6 +14,7 @@ import {
   UIChatMessage,
   UpdateMessageParams,
   UpdateMessageRAGParams,
+  UpdateMessageResult,
 } from '@lobechat/types';
 import type { HeatmapsProps } from '@lobehub/charts';
 import dayjs from 'dayjs';
@@ -54,6 +55,7 @@ export class MessageModel {
   query = async (
     { current = 0, pageSize = 1000, sessionId, topicId, groupId }: QueryMessageParams = {},
     options: {
+      groupAssistantMessages?: boolean;
       postProcessUrl?: (path: string | null, file: { fileType: string }) => Promise<string>;
     } = {},
   ) => {
@@ -211,7 +213,7 @@ export class MessageModel {
       .from(messageQueries)
       .where(inArray(messageQueries.messageId, messageIds));
 
-    return result.map(
+    const mappedMessages = result.map(
       ({ model, provider, translate, ttsId, ttsFile, ttsContentMd5, ttsVoice, ...item }) => {
         const messageQuery = messageQueriesList.find((relation) => relation.messageId === item.id);
         return {
@@ -224,8 +226,8 @@ export class MessageModel {
             })),
 
           extra: {
-            fromModel: model,
-            fromProvider: provider,
+            model: model,
+            provider: provider,
             translate,
             tts: ttsId
               ? {
@@ -246,13 +248,15 @@ export class MessageModel {
               size: size!,
               url,
             })),
-
           imageList: imageList
             .filter((relation) => relation.messageId === item.id)
             // eslint-disable-next-line @typescript-eslint/no-unused-vars
             .map<ChatImageItem>(({ id, url, name }) => ({ alt: name!, id, url })),
-
           meta: {},
+
+          model,
+
+          provider,
           ragQuery: messageQuery?.rewriteQuery,
           ragQueryId: messageQuery?.id,
           ragRawQuery: messageQuery?.userQuery,
@@ -263,6 +267,8 @@ export class MessageModel {
         } as unknown as UIChatMessage;
       },
     );
+
+    return mappedMessages;
   };
 
   findById = async (id: string) => {
@@ -455,8 +461,8 @@ export class MessageModel {
 
   create = async (
     {
-      fromModel,
-      fromProvider,
+      model: fromModel,
+      provider: fromProvider,
       files,
       plugin,
       pluginState,
@@ -522,54 +528,6 @@ export class MessageModel {
     });
   };
 
-  /**
-   * Create a new message and return the complete message list
-   *
-   * This method combines message creation and querying into a single operation,
-   * reducing the need for separate refresh calls and improving performance.
-   *
-   * @param params - Message creation parameters
-   * @param options - Query options for post-processing
-   * @returns Object containing the created message ID and full message list
-   *
-   * @example
-   * const { id, messages } = await messageModel.createNewMessage({
-   *   role: 'assistant',
-   *   content: 'Hello',
-   *   tools: [...],
-   *   sessionId: 'session-1',
-   * });
-   * // messages already contains grouped structure, no need to refresh
-   */
-  createNewMessage = async (
-    params: CreateMessageParams,
-    options: {
-      postProcessUrl?: (path: string | null, file: { fileType: string }) => Promise<string>;
-    } = {},
-  ): Promise<CreateMessageResult> => {
-    // 1. Create the message (reuse existing create method)
-    const item = await this.create(params);
-
-    // 2. Query all messages for this session/topic
-    // query() method internally applies groupAssistantMessages transformation
-    const messages = await this.query(
-      {
-        current: 0,
-        groupId: params.groupId,
-        pageSize: 9999,
-        sessionId: params.sessionId,
-        topicId: params.topicId, // Get all messages
-      },
-      options,
-    );
-
-    // 3. Return the result
-    return {
-      id: item.id,
-      messages,
-    };
-  };
-
   batchCreate = async (newMessages: DBMessageItem[]) => {
     const messagesToInsert = newMessages.map((m) => {
       // TODO: need a better way to handle this
@@ -589,27 +547,54 @@ export class MessageModel {
   };
   // **************** Update *************** //
 
-  update = async (id: string, { imageList, ...message }: Partial<UpdateMessageParams>) => {
-    return this.db.transaction(async (trx) => {
-      // 1. insert message files
-      if (imageList && imageList.length > 0) {
+  update = async (
+    id: string,
+    { imageList, ...message }: Partial<UpdateMessageParams>,
+    options?: {
+      groupAssistantMessages?: boolean;
+      postProcessUrl?: (path: string | null, file: { fileType: string }) => Promise<string>;
+      sessionId?: string | null;
+      topicId?: string | null;
+    },
+  ): Promise<UpdateMessageResult> => {
+    try {
+      await this.db.transaction(async (trx) => {
+        // 1. insert message files
+        if (imageList && imageList.length > 0) {
+          await trx
+            .insert(messagesFiles)
+            .values(
+              imageList.map((file) => ({ fileId: file.id, messageId: id, userId: this.userId })),
+            );
+        }
+
         await trx
-          .insert(messagesFiles)
-          .values(
-            imageList.map((file) => ({ fileId: file.id, messageId: id, userId: this.userId })),
-          );
+          .update(messages)
+          .set({ ...message })
+          .where(and(eq(messages.id, id), eq(messages.userId, this.userId)));
+      });
+
+      // if sessionId or topicId provided, return the updated message list
+      if (options?.sessionId !== undefined || options?.topicId !== undefined) {
+        const messageList = await this.query(
+          {
+            sessionId: options.sessionId,
+            topicId: options.topicId,
+          },
+          {
+            groupAssistantMessages: options.groupAssistantMessages ?? false,
+            postProcessUrl: options.postProcessUrl,
+          },
+        );
+
+        return { messages: messageList, success: true };
       }
 
-      return trx
-        .update(messages)
-        .set({
-          ...message,
-          // TODO: need a better way to handle this
-          // TODO: but I forget why 🤡
-          role: message.role as any,
-        })
-        .where(and(eq(messages.id, id), eq(messages.userId, this.userId)));
-    });
+      return { success: true };
+    } catch (error) {
+      console.error('Update message error:', error);
+      return { success: false };
+    }
   };
 
   updateMetadata = async (id: string, metadata: Record<string, any>) => {
@@ -625,16 +610,41 @@ export class MessageModel {
       .where(and(eq(messages.userId, this.userId), eq(messages.id, id)));
   };
 
-  updatePluginState = async (id: string, state: Record<string, any>) => {
+  updatePluginState = async (
+    id: string,
+    state: Record<string, any>,
+    options?: {
+      groupAssistantMessages?: boolean;
+      postProcessUrl?: (path: string | null, file: { fileType: string }) => Promise<string>;
+      sessionId?: string | null;
+      topicId?: string | null;
+    },
+  ): Promise<UpdateMessageResult> => {
     const item = await this.db.query.messagePlugins.findFirst({
       where: eq(messagePlugins.id, id),
     });
     if (!item) throw new Error('Plugin not found');
 
-    return this.db
+    await this.db
       .update(messagePlugins)
       .set({ state: merge(item.state || {}, state) })
       .where(eq(messagePlugins.id, id));
+
+    // Return updated messages if sessionId or topicId is provided
+    if (options?.sessionId !== undefined || options?.topicId !== undefined) {
+      const messageList = await this.query(
+        {
+          sessionId: options.sessionId,
+          topicId: options.topicId,
+        },
+        {
+          groupAssistantMessages: options.groupAssistantMessages ?? false,
+          postProcessUrl: options.postProcessUrl,
+        },
+      );
+      return { messages: messageList, success: true };
+    }
+    return { success: true };
   };
 
   updateMessagePlugin = async (id: string, value: Partial<MessagePluginItem>) => {
@@ -778,8 +788,11 @@ export class MessageModel {
 
   private genId = () => idGenerator('messages', 14);
 
-  private matchSession = (sessionId?: string | null) =>
-    sessionId ? eq(messages.sessionId, sessionId) : isNull(messages.sessionId);
+  private matchSession = (sessionId?: string | null) => {
+    if (sessionId === INBOX_SESSION_ID) return isNull(messages.sessionId);
+
+    return sessionId ? eq(messages.sessionId, sessionId) : isNull(messages.sessionId);
+  };
 
   private matchTopic = (topicId?: string | null) =>
     topicId ? eq(messages.topicId, topicId) : isNull(messages.topicId);

package/packages/database/src/models/session.ts

@@ -53,13 +53,44 @@ export class SessionModel {
   query = async ({ current = 0, pageSize = 9999 } = {}) => {
     const offset = current * pageSize;
 
-    return this.db.query.sessions.findMany({
-      limit: pageSize,
-      offset,
-      orderBy: [desc(sessions.updatedAt)],
-      where: and(eq(sessions.userId, this.userId), not(eq(sessions.slug, INBOX_SESSION_ID))),
-      with: { agentsToSessions: { columns: {}, with: { agent: true } }, group: true },
-    });
+    // Use leftJoin instead of nested with for better performance
+    const result = await this.db
+      .select({
+        // Agent fields (from agentsToSessions join)
+        agent: agents,
+        // Group fields
+        group: sessionGroups,
+        // Session fields
+        session: sessions,
+      })
+      .from(sessions)
+      .leftJoin(agentsToSessions, eq(sessions.id, agentsToSessions.sessionId))
+      .leftJoin(agents, eq(agentsToSessions.agentId, agents.id))
+      .leftJoin(sessionGroups, eq(sessions.groupId, sessionGroups.id))
+      .where(and(eq(sessions.userId, this.userId), not(eq(sessions.slug, INBOX_SESSION_ID))))
+      .orderBy(desc(sessions.updatedAt))
+      .limit(pageSize)
+      .offset(offset);
+
+    // Group results by session (since leftJoin can create multiple rows per session)
+    // Use Map to preserve order
+    const groupedResults = new Map<string, any>();
+
+    for (const row of result) {
+      const sessionId = row.session.id;
+      if (!groupedResults.has(sessionId)) {
+        groupedResults.set(sessionId, {
+          ...row.session,
+          agentsToSessions: [],
+          group: row.group,
+        });
+      }
+      if (row.agent) {
+        groupedResults.get(sessionId)!.agentsToSessions.push({ agent: row.agent });
+      }
+    }
+
+    return Array.from(groupedResults.values());
   };
 
   queryWithGroups = async (): Promise<ChatSessionList> => {
@@ -71,9 +102,11 @@ export class SessionModel {
       where: eq(sessions.userId, this.userId),
     });
 
+    const mappedSessions = result.map((item) => this.mapSessionItem(item as any));
+
     return {
       sessionGroups: groups as unknown as ChatSessionList['sessionGroups'],
-      sessions: result.map((item) => this.mapSessionItem(item as any)),
+      sessions: mappedSessions,
     };
   };
 
@@ -90,17 +123,28 @@ export class SessionModel {
   findByIdOrSlug = async (
     idOrSlug: string,
   ): Promise<(SessionItem & { agent: AgentItem }) | undefined> => {
-    const result = await this.db.query.sessions.findFirst({
-      where: and(
-        or(eq(sessions.id, idOrSlug), eq(sessions.slug, idOrSlug)),
-        eq(sessions.userId, this.userId),
-      ),
-      with: { agentsToSessions: { columns: {}, with: { agent: true } }, group: true },
-    });
+    // Use leftJoin instead of nested 'with' for better performance
+    const result = await this.db
+      .select({
+        agent: agents,
+        group: sessionGroups,
+        session: sessions,
+      })
+      .from(sessions)
+      .where(
+        and(
+          or(eq(sessions.id, idOrSlug), eq(sessions.slug, idOrSlug)),
+          eq(sessions.userId, this.userId),
+        ),
+      )
+      .leftJoin(agentsToSessions, eq(sessions.id, agentsToSessions.sessionId))
+      .leftJoin(agents, eq(agentsToSessions.agentId, agents.id))
+      .leftJoin(sessionGroups, eq(sessions.groupId, sessionGroups.id))
+      .limit(1);
 
-    if (!result) return;
+    if (!result || !result[0]) return;
 
-    return { ...result, agent: (result?.agentsToSessions?.[0] as any)?.agent } as any;
+    return { ...result[0].session, agent: result[0].agent, group: result[0].group } as any;
   };
 
   count = async (params?: {
@@ -332,7 +376,7 @@ export class SessionModel {
         .delete(agentsToSessions)
         .where(and(eq(agentsToSessions.sessionId, id), eq(agentsToSessions.userId, this.userId)));
 
-      // Delete the session
+      // Delete the session (this will cascade delete messages, topics, etc.)
       const result = await trx
         .delete(sessions)
         .where(and(eq(sessions.id, id), eq(sessions.userId, this.userId)));
@@ -397,17 +441,25 @@ export class SessionModel {
   };
 
   clearOrphanAgent = async (agentIds: string[], trx: any) => {
-    // Delete orphaned agents (those not linked to any other sessions)
-    for (const agentId of agentIds) {
-      const remaining = await trx
-        .select()
-        .from(agentsToSessions)
-        .where(eq(agentsToSessions.agentId, agentId))
-        .limit(1);
+    if (agentIds.length === 0) return;
 
-      if (remaining.length === 0) {
-        await trx.delete(agents).where(and(eq(agents.id, agentId), eq(agents.userId, this.userId)));
-      }
+    // Batch query to find which agents still have sessions
+    const remainingLinks = (await trx
+      .select({ agentId: agentsToSessions.agentId })
+      .from(agentsToSessions)
+      .where(inArray(agentsToSessions.agentId, agentIds))) as { agentId: string }[];
+
+    const linkedAgentIds = new Set(remainingLinks.map((link) => link.agentId));
+
+    // Find orphaned agents (those not in the linked set)
+    const orphanedAgentIds = agentIds.filter((id) => !linkedAgentIds.has(id));
+
+    // Batch delete orphaned agents (this will cascade to agentsFiles, agentsKnowledgeBases, etc.)
+    // and SET NULL on messages.agentId
+    if (orphanedAgentIds.length > 0) {
+      await trx
+        .delete(agents)
+        .where(and(inArray(agents.id, orphanedAgentIds), eq(agents.userId, this.userId)));
     }
   };
 

package/packages/database/src/models/user.ts

@@ -82,7 +82,8 @@ export class UserModel {
       })
       .from(users)
       .where(eq(users.id, this.userId))
-      .leftJoin(userSettings, eq(users.id, userSettings.id));
+      .leftJoin(userSettings, eq(users.id, userSettings.id))
+      .limit(1);
 
     if (!result || !result[0]) {
       throw new UserNotFoundError();

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

@@ -731,6 +731,105 @@ describe('AiInfraRepos', () => {
       // 无 settings
       expect(merged?.settings).toBeUndefined();
     });
+
+    it('should prefer user settings over builtin settings', async () => {
+      const mockProviders = [
+        { enabled: true, id: 'openai', name: 'OpenAI', source: 'builtin' as const },
+      ];
+
+      const userModel: EnabledAiModel = {
+        id: 'gpt-4',
+        providerId: 'openai',
+        enabled: true,
+        type: 'chat',
+        abilities: {},
+        settings: { searchImpl: 'params', searchProvider: 'user-provider' },
+      };
+
+      const builtinModel = {
+        id: 'gpt-4',
+        enabled: true,
+        type: 'chat' as const,
+        settings: { searchImpl: 'tool', searchProvider: 'builtin-provider' },
+      };
+
+      vi.spyOn(repo, 'getAiProviderList').mockResolvedValue(mockProviders);
+      vi.spyOn(repo.aiModelModel, 'getAllModels').mockResolvedValue([userModel]);
+      vi.spyOn(repo as any, 'fetchBuiltinModels').mockResolvedValue([builtinModel]);
+
+      const result = await repo.getEnabledModels();
+
+      const merged = result.find((m) => m.id === 'gpt-4');
+      expect(merged).toBeDefined();
+      // 应该使用用户的 settings，不是内置的
+      expect(merged?.settings).toEqual({ searchImpl: 'params', searchProvider: 'user-provider' });
+    });
+
+    it('should use builtin settings when user has no settings', async () => {
+      const mockProviders = [
+        { enabled: true, id: 'openai', name: 'OpenAI', source: 'builtin' as const },
+      ];
+
+      const userModel: EnabledAiModel = {
+        id: 'gpt-4',
+        providerId: 'openai',
+        enabled: true,
+        type: 'chat',
+        abilities: { vision: true },
+        // 用户未设置 settings
+      };
+
+      const builtinModel = {
+        id: 'gpt-4',
+        enabled: true,
+        type: 'chat' as const,
+        settings: { searchImpl: 'tool', searchProvider: 'google' },
+      };
+
+      vi.spyOn(repo, 'getAiProviderList').mockResolvedValue(mockProviders);
+      vi.spyOn(repo.aiModelModel, 'getAllModels').mockResolvedValue([userModel]);
+      vi.spyOn(repo as any, 'fetchBuiltinModels').mockResolvedValue([builtinModel]);
+
+      const result = await repo.getEnabledModels();
+
+      const merged = result.find((m) => m.id === 'gpt-4');
+      expect(merged).toBeDefined();
+      // 应该使用内置的 settings
+      expect(merged?.settings).toEqual({ searchImpl: 'tool', searchProvider: 'google' });
+    });
+
+    it('should have no settings when both user and builtin have no settings', async () => {
+      const mockProviders = [
+        { enabled: true, id: 'openai', name: 'OpenAI', source: 'builtin' as const },
+      ];
+
+      const userModel: EnabledAiModel = {
+        id: 'gpt-4',
+        providerId: 'openai',
+        enabled: true,
+        type: 'chat',
+        abilities: { vision: true },
+        // 用户未设置 settings
+      };
+
+      const builtinModel = {
+        id: 'gpt-4',
+        enabled: true,
+        type: 'chat' as const,
+        // 内置也无 settings
+      };
+
+      vi.spyOn(repo, 'getAiProviderList').mockResolvedValue(mockProviders);
+      vi.spyOn(repo.aiModelModel, 'getAllModels').mockResolvedValue([userModel]);
+      vi.spyOn(repo as any, 'fetchBuiltinModels').mockResolvedValue([builtinModel]);
+
+      const result = await repo.getEnabledModels();
+
+      const merged = result.find((m) => m.id === 'gpt-4');
+      expect(merged).toBeDefined();
+      // 无 settings
+      expect(merged?.settings).toBeUndefined();
+    });
   });
 
   describe('getAiProviderModelList', () => {
@@ -1270,6 +1369,105 @@ describe('AiInfraRepos', () => {
       // 无 settings
       expect(merged?.settings).toBeUndefined();
     });
+
+    it('should prefer user settings over builtin settings in getAiProviderModelList', async () => {
+      const providerId = 'openai';
+
+      const userModels: AiProviderModelListItem[] = [
+        {
+          id: 'gpt-4',
+          type: 'chat',
+          enabled: true,
+          abilities: {},
+          settings: { searchImpl: 'params', searchProvider: 'user-provider' },
+        } as any,
+      ];
+
+      const builtinModels: AiProviderModelListItem[] = [
+        {
+          id: 'gpt-4',
+          type: 'chat',
+          enabled: true,
+          settings: { searchImpl: 'tool', searchProvider: 'builtin-provider' },
+        } as any,
+      ];
+
+      vi.spyOn(repo.aiModelModel, 'getModelListByProviderId').mockResolvedValue(userModels);
+      vi.spyOn(repo as any, 'fetchBuiltinModels').mockResolvedValue(builtinModels);
+
+      const result = await repo.getAiProviderModelList(providerId);
+
+      const merged = result.find((m) => m.id === 'gpt-4');
+      expect(merged).toBeDefined();
+      // 应该使用用户的 settings
+      expect(merged?.settings).toEqual({ searchImpl: 'params', searchProvider: 'user-provider' });
+    });
+
+    it('should use builtin settings when user has no settings in getAiProviderModelList', async () => {
+      const providerId = 'openai';
+
+      const userModels: AiProviderModelListItem[] = [
+        {
+          id: 'gpt-4',
+          type: 'chat',
+          enabled: true,
+          abilities: { vision: true },
+          // 用户未设置 settings
+        },
+      ];
+
+      const builtinModels: AiProviderModelListItem[] = [
+        {
+          id: 'gpt-4',
+          type: 'chat',
+          enabled: true,
+          settings: { searchImpl: 'tool', searchProvider: 'google' },
+        } as any,
+      ];
+
+      vi.spyOn(repo.aiModelModel, 'getModelListByProviderId').mockResolvedValue(userModels);
+      vi.spyOn(repo as any, 'fetchBuiltinModels').mockResolvedValue(builtinModels);
+
+      const result = await repo.getAiProviderModelList(providerId);
+
+      const merged = result.find((m) => m.id === 'gpt-4');
+      expect(merged).toBeDefined();
+      // 应该使用内置的 settings
+      expect(merged?.settings).toEqual({ searchImpl: 'tool', searchProvider: 'google' });
+    });
+
+    it('should have no settings when both user and builtin have no settings in getAiProviderModelList', async () => {
+      const providerId = 'openai';
+
+      const userModels: AiProviderModelListItem[] = [
+        {
+          id: 'gpt-4',
+          type: 'chat',
+          enabled: true,
+          abilities: { vision: true },
+          // 用户未设置 settings
+        },
+      ];
+
+      const builtinModels: AiProviderModelListItem[] = [
+        {
+          id: 'gpt-4',
+          type: 'chat',
+          enabled: true,
+          // 内置也无 settings
+        },
+      ];
+
+      vi.spyOn(repo.aiModelModel, 'getModelListByProviderId').mockResolvedValue(userModels);
+      vi.spyOn(repo as any, 'fetchBuiltinModels').mockResolvedValue(builtinModels);
+
+      const result = await repo.getAiProviderModelList(providerId);
+
+      const merged = result.find((m) => m.id === 'gpt-4');
+      expect(merged).toBeDefined();
+      // 无 settings
+      expect(merged?.settings).toBeUndefined();
+    });
   });
 
   describe('getAiProviderRuntimeState', () => {