@lobehub/lobehub 2.0.0-next.2 → 2.0.0-next.20

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

@@ -1,3 +1,4 @@
+import { INBOX_SESSION_ID } from '@lobechat/const';
 import {
   ChatFileItem,
   ChatImageItem,
@@ -14,6 +15,7 @@ import {
   UIChatMessage,
   UpdateMessageParams,
   UpdateMessageRAGParams,
+  UpdateMessageResult,
 } from '@lobechat/types';
 import type { HeatmapsProps } from '@lobehub/charts';
 import dayjs from 'dayjs';
@@ -39,6 +41,7 @@ import {
 } from '../schemas';
 import { LobeChatDatabase } from '../type';
 import { genEndDateWhere, genRangeWhere, genStartDateWhere, genWhere } from '../utils/genWhere';
+import { groupAssistantMessages } from '../utils/groupMessages';
 import { idGenerator } from '../utils/idGenerator';
 
 export class MessageModel {
@@ -54,6 +57,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 +215,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 {
@@ -246,13 +250,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 +269,10 @@ export class MessageModel {
         } as unknown as UIChatMessage;
       },
     );
+
+    // Group assistant messages with their tool results
+    const { groupAssistantMessages: useGroup = false } = options;
+    return useGroup ? groupAssistantMessages(mappedMessages) : mappedMessages;
   };
 
   findById = async (id: string) => {
@@ -560,7 +570,7 @@ export class MessageModel {
         sessionId: params.sessionId,
         topicId: params.topicId, // Get all messages
       },
-      options,
+      { ...options, groupAssistantMessages: true },
     );
 
     // 3. Return the result
@@ -589,27 +599,52 @@ 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?: {
+      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,
+          },
+          {
+            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>) => {
@@ -778,8 +813,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,
     };
   };
 
@@ -332,7 +365,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 +430,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/repositories/dataImporter/deprecated/__tests__/index.test.ts

@@ -12,12 +12,13 @@ import {
   topics,
   users,
 } from '@/database/schemas';
-import { CURRENT_CONFIG_VERSION } from '@/migrations';
 import { ImporterEntryData } from '@/types/importer';
 
 import { DeprecatedDataImporterRepos as DataImporterRepos } from '../index';
 import mockImportData from './fixtures/messages.json';
 
+const CURRENT_CONFIG_VERSION = 7;
+
 const serverDB = await getTestDBInstance();
 
 const userId = 'test-user-id';

package/packages/database/src/repositories/dataImporter/deprecated/index.ts

@@ -1,6 +1,5 @@
 import { and, eq, inArray, sql } from 'drizzle-orm';
 
-import { ImportResult } from '@/services/import/_deprecated';
 import { ImporterEntryData } from '@/types/importer';
 import { sanitizeUTF8 } from '@/utils/sanitizeUTF8';
 
@@ -16,6 +15,13 @@ import {
 } from '../../../schemas';
 import { LobeChatDatabase } from '../../../type';
 
+interface ImportResult {
+  added: number;
+  errors: number;
+  skips: number;
+  updated?: number;
+}
+
 export class DeprecatedDataImporterRepos {
   private userId: string;
   private db: LobeChatDatabase;

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

@@ -62,11 +62,11 @@ export const agents = pgTable(
 
     ...timestamps,
   },
-  (t) => ({
-    clientIdUnique: uniqueIndex('client_id_user_id_unique').on(t.clientId, t.userId),
-    titleIndex: index('agents_title_idx').on(t.title),
-    descriptionIndex: index('agents_description_idx').on(t.description),
-  }),
+  (t) => [
+    uniqueIndex('client_id_user_id_unique').on(t.clientId, t.userId),
+    index('agents_title_idx').on(t.title),
+    index('agents_description_idx').on(t.description),
+  ],
 );
 
 export const insertAgentSchema = createInsertSchema(agents);
@@ -90,9 +90,7 @@ export const agentsKnowledgeBases = pgTable(
 
     ...timestamps,
   },
-  (t) => ({
-    pk: primaryKey({ columns: [t.agentId, t.knowledgeBaseId] }),
-  }),
+  (t) => [primaryKey({ columns: [t.agentId, t.knowledgeBaseId] })],
 );
 
 export const agentsFiles = pgTable(
@@ -111,7 +109,8 @@ export const agentsFiles = pgTable(
 
     ...timestamps,
   },
-  (t) => ({
-    pk: primaryKey({ columns: [t.fileId, t.agentId, t.userId] }),
-  }),
+  (t) => [
+    primaryKey({ columns: [t.fileId, t.agentId, t.userId] }),
+    index('agents_files_agent_id_idx').on(t.agentId),
+  ],
 );

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

@@ -62,7 +62,10 @@ export const messageGroups = pgTable(
 
     ...timestamps,
   },
-  (t) => [uniqueIndex('message_groups_client_id_user_id_unique').on(t.clientId, t.userId)],
+  (t) => [
+    uniqueIndex('message_groups_client_id_user_id_unique').on(t.clientId, t.userId),
+    index('message_groups_topic_id_idx').on(t.topicId),
+  ],
 );
 
 export const insertMessageGroupSchema = createInsertSchema(messageGroups);
@@ -130,6 +133,7 @@ export const messages = pgTable(
     index('messages_user_id_idx').on(table.userId),
     index('messages_session_id_idx').on(table.sessionId),
     index('messages_thread_id_idx').on(table.threadId),
+    index('messages_agent_id_idx').on(table.agentId),
   ],
 );
 

package/packages/database/src/schemas/relations.ts

@@ -1,6 +1,6 @@
 /* eslint-disable sort-keys-fix/sort-keys-fix  */
 import { relations } from 'drizzle-orm';
-import { pgTable, primaryKey, text, uuid, varchar } from 'drizzle-orm/pg-core';
+import { index, pgTable, primaryKey, text, uuid, varchar } from 'drizzle-orm/pg-core';
 
 import { createdAt } from './_helpers';
 import { agents, agentsFiles, agentsKnowledgeBases } from './agent';
@@ -28,9 +28,11 @@ export const agentsToSessions = pgTable(
       .references(() => users.id, { onDelete: 'cascade' })
       .notNull(),
   },
-  (t) => ({
-    pk: primaryKey({ columns: [t.agentId, t.sessionId] }),
-  }),
+  (t) => [
+    primaryKey({ columns: [t.agentId, t.sessionId] }),
+    index('agents_to_sessions_session_id_idx').on(t.sessionId),
+    index('agents_to_sessions_agent_id_idx').on(t.agentId),
+  ],
 );
 
 export const filesToSessions = pgTable(

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

@@ -70,6 +70,8 @@ export const sessions = pgTable(
 
     index('sessions_user_id_idx').on(t.userId),
     index('sessions_id_user_id_idx').on(t.id, t.userId),
+    index('sessions_user_id_updated_at_idx').on(t.userId, t.updatedAt),
+    index('sessions_group_id_idx').on(t.groupId),
   ],
 );
 

package/packages/database/src/schemas/topic.ts

@@ -33,6 +33,8 @@ export const topics = pgTable(
     uniqueIndex('topics_client_id_user_id_unique').on(t.clientId, t.userId),
     index('topics_user_id_idx').on(t.userId),
     index('topics_id_user_id_idx').on(t.id, t.userId),
+    index('topics_session_id_idx').on(t.sessionId),
+    index('topics_group_id_idx').on(t.groupId),
   ],
 );
 
@@ -65,7 +67,10 @@ export const threads = pgTable(
     lastActiveAt: timestamptz('last_active_at').defaultNow(),
     ...timestamps,
   },
-  (t) => [uniqueIndex('threads_client_id_user_id_unique').on(t.clientId, t.userId)],
+  (t) => [
+    uniqueIndex('threads_client_id_user_id_unique').on(t.clientId, t.userId),
+    index('threads_topic_id_idx').on(t.topicId),
+  ],
 );
 
 export type NewThread = typeof threads.$inferInsert;

package/packages/database/src/utils/__tests__/groupMessages.test.ts

@@ -54,6 +54,7 @@ describe('groupAssistantMessages', () => {
           content: 'Beijing: Sunny, 25°C',
           state: { cached: true },
         },
+        result_msg_id: 'msg-2',
       });
     });
 
@@ -112,7 +113,9 @@ describe('groupAssistantMessages', () => {
       const block = result[0].children![0];
       expect(block.tools).toHaveLength(2);
       expect(block.tools![0].result?.content).toBe('Beijing: Sunny, 25°C');
+      expect(block.tools![0].result_msg_id).toBe('msg-2');
       expect(block.tools![1].result?.content).toBe('Latest tech news: AI breakthrough');
+      expect(block.tools![1].result_msg_id).toBe('msg-3');
     });
 
     it('should handle assistant message without tools', () => {
@@ -165,6 +168,7 @@ describe('groupAssistantMessages', () => {
       const block = result[0].children![0];
       expect(block.tools).toHaveLength(1);
       expect(block.tools![0].result).toBeUndefined();
+      expect(block.tools![0].result_msg_id).toBeUndefined();
     });
   });
 
@@ -366,11 +370,13 @@ describe('groupAssistantMessages', () => {
       // First child: original assistant with tool result
       expect(result[0].children![0].id).toBe('msg-1');
       expect(result[0].children![0].tools![0].result?.content).toBe('Sunny, 25°C');
+      expect(result[0].children![0].tools![0].result_msg_id).toBe('msg-2');
 
       // Second child: follow-up assistant with its own tool result
       expect(result[0].children![1].id).toBe('msg-3');
       expect(result[0].children![1].tools).toHaveLength(1);
       expect(result[0].children![1].tools![0].result?.content).toBe('Breaking news');
+      expect(result[0].children![1].tools![0].result_msg_id).toBe('msg-4');
     });
 
     it('should group multiple follow-up assistants in chain (3+ assistants)', () => {
@@ -448,9 +454,11 @@ describe('groupAssistantMessages', () => {
 
       expect(result[0].children![0].id).toBe('msg-1');
       expect(result[0].children![0].tools![0].result?.content).toBe('Result 1');
+      expect(result[0].children![0].tools![0].result_msg_id).toBe('msg-2');
 
       expect(result[0].children![1].id).toBe('msg-3');
       expect(result[0].children![1].tools![0].result?.content).toBe('Result 2');
+      expect(result[0].children![1].tools![0].result_msg_id).toBe('msg-4');
 
       expect(result[0].children![2].id).toBe('msg-5');
       expect(result[0].children![2].content).toBe('Step 3 final');
@@ -677,7 +685,6 @@ describe('groupAssistantMessages', () => {
       expect(block.content).toBe('Test');
       expect(block.tools).toHaveLength(1);
       expect(block.imageList).toHaveLength(1);
-      expect(block.fileList).toHaveLength(1);
     });
 
     it('should preserve all tool result fields', () => {
@@ -720,6 +727,7 @@ describe('groupAssistantMessages', () => {
         state: { step: 1 },
         error: null,
       });
+      expect(block.tools![0].result_msg_id).toBe('msg-2');
     });
   });
 
@@ -928,6 +936,142 @@ describe('groupAssistantMessages', () => {
       expect(result[0].performance).toBeUndefined();
       expect(result[0].metadata).toBeUndefined();
     });
+
+    it('should map reasoning field to reasoning in children blocks', () => {
+      const input: UIChatMessage[] = [
+        {
+          id: 'msg-1',
+          role: 'assistant',
+          content: 'Test response',
+          reasoning: {
+            content: 'This is my reasoning process',
+            duration: 1500,
+          },
+          tools: [
+            {
+              id: 'tool-1',
+              identifier: 'test',
+              apiName: 'test',
+              arguments: '{}',
+              type: 'default',
+            },
+          ],
+          createdAt: Date.now(),
+          updatedAt: Date.now(),
+          meta: {},
+        } as UIChatMessage,
+        {
+          id: 'msg-2',
+          role: 'tool',
+          tool_call_id: 'tool-1',
+          content: 'Tool result',
+          createdAt: Date.now(),
+          updatedAt: Date.now(),
+          meta: {},
+        } as UIChatMessage,
+        {
+          id: 'msg-3',
+          role: 'assistant',
+          content: 'Follow-up response',
+          parentId: 'msg-2',
+          reasoning: {
+            content: 'Follow-up reasoning',
+            duration: 2000,
+          },
+          createdAt: Date.now(),
+          updatedAt: Date.now(),
+          meta: {},
+        } as UIChatMessage,
+      ];
+
+      const result = groupAssistantMessages(input);
+
+      // First block should have reasoning
+      expect(result[0].children![0].reasoning).toEqual({
+        content: 'This is my reasoning process',
+        duration: 1500,
+      });
+
+      // Second block (follow-up) should also have reasoning
+      expect(result[0].children![1].reasoning).toEqual({
+        content: 'Follow-up reasoning',
+        duration: 2000,
+      });
+
+      // Group message should not have reasoning (moved to children)
+      expect(result[0].reasoning).toBeUndefined();
+    });
+
+    it('should preserve error field in children blocks', () => {
+      const input = [
+        {
+          id: 'msg-1',
+          role: 'assistant',
+          content: 'Failed to process',
+          error: {
+            type: 'InvalidAPIKey',
+            message: 'API key is invalid',
+          },
+          tools: [
+            {
+              id: 'tool-1',
+              identifier: 'test',
+              apiName: 'test',
+              arguments: '{}',
+              type: 'default',
+            },
+          ],
+          createdAt: Date.now(),
+          updatedAt: Date.now(),
+          meta: {},
+        } as unknown as UIChatMessage,
+      ] as UIChatMessage[];
+
+      const result = groupAssistantMessages(input);
+
+      // Child block should have error
+      expect(result[0].children![0].error).toEqual({
+        type: 'InvalidAPIKey',
+        message: 'API key is invalid',
+      });
+    });
+
+    it('should preserve imageList in children blocks', () => {
+      const input: UIChatMessage[] = [
+        {
+          id: 'msg-1',
+          role: 'assistant',
+          content: 'Here are the images',
+          imageList: [
+            { id: 'img-1', url: 'https://example.com/img1.jpg', alt: 'Image 1' },
+            { id: 'img-2', url: 'https://example.com/img2.jpg', alt: 'Image 2' },
+          ],
+          tools: [
+            {
+              id: 'tool-1',
+              identifier: 'test',
+              apiName: 'test',
+              arguments: '{}',
+              type: 'default',
+            },
+          ],
+          createdAt: Date.now(),
+          updatedAt: Date.now(),
+          meta: {},
+        } as UIChatMessage,
+      ];
+
+      const result = groupAssistantMessages(input);
+
+      // Child block should have imageList
+      expect(result[0].children![0].imageList).toEqual([
+        { id: 'img-1', url: 'https://example.com/img1.jpg', alt: 'Image 1' },
+        { id: 'img-2', url: 'https://example.com/img2.jpg', alt: 'Image 2' },
+      ]);
+
+      // Parent should not have imageList (moved to children)
+      expect(result[0].imageList).toBeUndefined();
+    });
   });
 
   describe('Empty and Null Cases', () => {
@@ -958,7 +1102,6 @@ describe('groupAssistantMessages', () => {
 
       // Empty arrays should become undefined
       expect(result[0].children![0].imageList).toBeUndefined();
-      expect(result[0].children![0].fileList).toBeUndefined();
     });
 
     it('should handle empty message list', () => {