@lobehub/lobehub 2.0.0-next.29 → 2.0.0-next.30

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

@@ -11,6 +11,7 @@ import { getServerDB } from '@/database/server';
 import { authedProcedure, publicProcedure, router } from '@/libs/trpc/lambda';
 import { serverDatabase } from '@/libs/trpc/lambda/middleware';
 import { FileService } from '@/server/services/file';
+import { MessageService } from '@/server/services/message';
 
 const messageProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
   const { ctx } = opts;
@@ -19,6 +20,7 @@ const messageProcedure = authedProcedure.use(serverDatabase).use(async (opts) =>
     ctx: {
       fileService: new FileService(ctx.serverDB, ctx.userId),
       messageModel: new MessageModel(ctx.serverDB, ctx.userId),
+      messageService: new MessageService(ctx.serverDB, ctx.userId),
     },
   });
 });
@@ -53,11 +55,10 @@ export const messageRouter = router({
     }),
 
   createNewMessage: messageProcedure
-    .input(CreateNewMessageParamsSchema)
+    .input(CreateNewMessageParamsSchema.extend({ useGroup: z.boolean().optional() }))
     .mutation(async ({ input, ctx }) => {
-      return ctx.messageModel.createNewMessage(input as any, {
-        postProcessUrl: (path) => ctx.fileService.getFullFileUrl(path),
-      });
+      const { useGroup, ...params } = input;
+      return ctx.messageService.createNewMessage(params as any, { useGroup });
     }),
 
   getHeatmaps: messageProcedure.query(async ({ ctx }) => {
@@ -109,23 +110,8 @@ export const messageRouter = router({
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      await ctx.messageModel.deleteMessage(input.id);
-
-      // If sessionId or topicId is provided, return the full message list
-      if (input.sessionId !== undefined || input.topicId !== undefined) {
-        const messageList = await ctx.messageModel.query(
-          {
-            sessionId: input.sessionId,
-            topicId: input.topicId,
-          },
-          {
-            groupAssistantMessages: input.useGroup ?? false,
-            postProcessUrl: (path) => ctx.fileService.getFullFileUrl(path),
-          },
-        );
-        return { messages: messageList, success: true };
-      }
-      return { success: true };
+      const { id, ...options } = input;
+      return ctx.messageService.removeMessage(id, options);
     }),
 
   removeMessageQuery: messageProcedure
@@ -144,23 +130,8 @@ export const messageRouter = router({
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      await ctx.messageModel.deleteMessages(input.ids);
-
-      // If sessionId or topicId is provided, return the full message list
-      if (input.sessionId !== undefined || input.topicId !== undefined) {
-        const messageList = await ctx.messageModel.query(
-          {
-            sessionId: input.sessionId,
-            topicId: input.topicId,
-          },
-          {
-            groupAssistantMessages: input.useGroup ?? false,
-            postProcessUrl: (path) => ctx.fileService.getFullFileUrl(path),
-          },
-        );
-        return { messages: messageList, success: true };
-      }
-      return { success: true };
+      const { ids, ...options } = input;
+      return ctx.messageService.removeMessages(ids, options);
     }),
 
   removeMessagesByAssistant: messageProcedure
@@ -207,12 +178,8 @@ export const messageRouter = router({
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      return ctx.messageModel.update(input.id, input.value as any, {
-        groupAssistantMessages: input.useGroup ?? false,
-        postProcessUrl: (path) => ctx.fileService.getFullFileUrl(path),
-        sessionId: input.sessionId,
-        topicId: input.topicId,
-      });
+      const { id, value, ...options } = input;
+      return ctx.messageService.updateMessage(id, value as any, options);
     }),
 
   updateMessagePlugin: messageProcedure
@@ -235,23 +202,8 @@ export const messageRouter = router({
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      await ctx.messageModel.updateMessageRAG(input.id, input.value);
-
-      // If sessionId or topicId is provided, return the full message list
-      if (input.sessionId !== undefined || input.topicId !== undefined) {
-        const messageList = await ctx.messageModel.query(
-          {
-            sessionId: input.sessionId,
-            topicId: input.topicId,
-          },
-          {
-            groupAssistantMessages: input.useGroup ?? false,
-            postProcessUrl: (path) => ctx.fileService.getFullFileUrl(path),
-          },
-        );
-        return { messages: messageList, success: true };
-      }
-      return { success: true };
+      const { id, value, ...options } = input;
+      return ctx.messageService.updateMessageRAG(id, value, options);
     }),
 
   updateMetadata: messageProcedure
@@ -276,23 +228,8 @@ export const messageRouter = router({
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      // If sessionId or topicId is provided, we need to return the full message list
-      if (input.sessionId !== undefined || input.topicId !== undefined) {
-        await ctx.messageModel.updateMessagePlugin(input.id, { error: input.value });
-        const messageList = await ctx.messageModel.query(
-          {
-            sessionId: input.sessionId,
-            topicId: input.topicId,
-          },
-          {
-            groupAssistantMessages: input.useGroup ?? false,
-            postProcessUrl: (path) => ctx.fileService.getFullFileUrl(path),
-          },
-        );
-        return { messages: messageList, success: true };
-      }
-      await ctx.messageModel.updateMessagePlugin(input.id, { error: input.value });
-      return { success: true };
+      const { id, value, ...options } = input;
+      return ctx.messageService.updatePluginError(id, value, options);
     }),
 
   updatePluginState: messageProcedure
@@ -306,12 +243,8 @@ export const messageRouter = router({
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      return ctx.messageModel.updatePluginState(input.id, input.value, {
-        groupAssistantMessages: input.useGroup ?? false,
-        postProcessUrl: (path) => ctx.fileService.getFullFileUrl(path),
-        sessionId: input.sessionId,
-        topicId: input.topicId,
-      });
+      const { id, value, ...options } = input;
+      return ctx.messageService.updatePluginState(id, value, options);
     }),
 
   updateTTS: messageProcedure

package/src/server/services/message/__tests__/index.test.ts

@@ -0,0 +1,348 @@
+import { LobeChatDatabase } from '@lobechat/database';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { MessageModel } from '@/database/models/message';
+import { FileService } from '@/server/services/file';
+
+import { MessageService } from '../index';
+
+vi.mock('@/database/models/message');
+vi.mock('@/server/services/file');
+
+describe('MessageService', () => {
+  let messageService: MessageService;
+  let mockDB: LobeChatDatabase;
+  let mockMessageModel: MessageModel;
+  let mockFileService: FileService;
+  const userId = 'test-user-id';
+
+  beforeEach(() => {
+    mockDB = {} as LobeChatDatabase;
+    mockMessageModel = {
+      create: vi.fn(),
+      deleteMessage: vi.fn(),
+      deleteMessages: vi.fn(),
+      query: vi.fn(),
+      update: vi.fn(),
+      updateMessagePlugin: vi.fn(),
+      updateMessageRAG: vi.fn(),
+      updatePluginState: vi.fn(),
+    } as any;
+
+    mockFileService = {
+      getFullFileUrl: vi.fn().mockImplementation((path) => Promise.resolve(`/files${path}`)),
+    } as any;
+
+    // Mock constructors
+    vi.mocked(MessageModel).mockImplementation(() => mockMessageModel);
+    vi.mocked(FileService).mockImplementation(() => mockFileService);
+
+    messageService = new MessageService(mockDB, userId);
+  });
+
+  describe('removeMessage', () => {
+    it('should delete message and return { success: true } when no sessionId/topicId provided', async () => {
+      const messageId = 'msg-1';
+
+      const result = await messageService.removeMessage(messageId);
+
+      expect(mockMessageModel.deleteMessage).toHaveBeenCalledWith(messageId);
+      expect(result).toEqual({ success: true });
+      expect(mockMessageModel.query).not.toHaveBeenCalled();
+    });
+
+    it('should delete message and return message list when sessionId provided', async () => {
+      const messageId = 'msg-1';
+      const mockMessages = [{ id: 'msg-2', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.removeMessage(messageId, { sessionId: 'session-1' });
+
+      expect(mockMessageModel.deleteMessage).toHaveBeenCalledWith(messageId);
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        { groupId: undefined, sessionId: 'session-1', topicId: undefined },
+        expect.objectContaining({
+          groupAssistantMessages: false,
+        }),
+      );
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+
+    it('should delete message and return message list when topicId provided', async () => {
+      const messageId = 'msg-1';
+      const mockMessages = [{ id: 'msg-2', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.removeMessage(messageId, { topicId: 'topic-1' });
+
+      expect(mockMessageModel.deleteMessage).toHaveBeenCalledWith(messageId);
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        { groupId: undefined, sessionId: undefined, topicId: 'topic-1' },
+        expect.objectContaining({
+          groupAssistantMessages: false,
+        }),
+      );
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+  });
+
+  describe('removeMessages', () => {
+    it('should delete messages and return { success: true } when no sessionId/topicId provided', async () => {
+      const messageIds = ['msg-1', 'msg-2'];
+
+      const result = await messageService.removeMessages(messageIds);
+
+      expect(mockMessageModel.deleteMessages).toHaveBeenCalledWith(messageIds);
+      expect(result).toEqual({ success: true });
+      expect(mockMessageModel.query).not.toHaveBeenCalled();
+    });
+
+    it('should delete messages and return message list when sessionId provided', async () => {
+      const messageIds = ['msg-1', 'msg-2'];
+      const mockMessages = [{ id: 'msg-3', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.removeMessages(messageIds, { sessionId: 'session-1' });
+
+      expect(mockMessageModel.deleteMessages).toHaveBeenCalledWith(messageIds);
+      expect(mockMessageModel.query).toHaveBeenCalled();
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+  });
+
+  describe('updateMessageRAG', () => {
+    it('should update RAG and return { success: true } when no sessionId/topicId provided', async () => {
+      const messageId = 'msg-1';
+      const ragValue = { fileChunks: [{ id: 'chunk-1', similarity: 0.95 }] };
+
+      const result = await messageService.updateMessageRAG(messageId, ragValue);
+
+      expect(mockMessageModel.updateMessageRAG).toHaveBeenCalledWith(messageId, ragValue);
+      expect(result).toEqual({ success: true });
+      expect(mockMessageModel.query).not.toHaveBeenCalled();
+    });
+
+    it('should update RAG and return message list when sessionId provided', async () => {
+      const messageId = 'msg-1';
+      const ragValue = { fileChunks: [{ id: 'chunk-1', similarity: 0.95 }] };
+      const mockMessages = [{ id: 'msg-1', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.updateMessageRAG(messageId, ragValue, {
+        sessionId: 'session-1',
+      });
+
+      expect(mockMessageModel.updateMessageRAG).toHaveBeenCalledWith(messageId, ragValue);
+      expect(mockMessageModel.query).toHaveBeenCalled();
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+  });
+
+  describe('updatePluginError', () => {
+    it('should update plugin error and return { success: true } when no sessionId/topicId provided', async () => {
+      const messageId = 'msg-1';
+      const error = { type: 'TestError', message: 'Test error message' };
+
+      const result = await messageService.updatePluginError(messageId, error);
+
+      expect(mockMessageModel.updateMessagePlugin).toHaveBeenCalledWith(messageId, { error });
+      expect(result).toEqual({ success: true });
+      expect(mockMessageModel.query).not.toHaveBeenCalled();
+    });
+
+    it('should update plugin error and return message list when sessionId provided', async () => {
+      const messageId = 'msg-1';
+      const error = { type: 'TestError', message: 'Test error message' };
+      const mockMessages = [{ id: 'msg-1', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.updatePluginError(messageId, error, {
+        sessionId: 'session-1',
+      });
+
+      expect(mockMessageModel.updateMessagePlugin).toHaveBeenCalledWith(messageId, { error });
+      expect(mockMessageModel.query).toHaveBeenCalled();
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+  });
+
+  describe('updatePluginState', () => {
+    it('should update plugin state and return { success: true } when no sessionId/topicId provided', async () => {
+      const messageId = 'msg-1';
+      const state = { key: 'value' };
+
+      const result = await messageService.updatePluginState(messageId, state, {});
+
+      expect(mockMessageModel.updatePluginState).toHaveBeenCalledWith(messageId, state);
+      expect(result).toEqual({ success: true });
+      expect(mockMessageModel.query).not.toHaveBeenCalled();
+    });
+
+    it('should update plugin state and return message list when sessionId provided', async () => {
+      const messageId = 'msg-1';
+      const state = { key: 'value' };
+      const mockMessages = [{ id: 'msg-1', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.updatePluginState(messageId, state, {
+        sessionId: 'session-1',
+      });
+
+      expect(mockMessageModel.updatePluginState).toHaveBeenCalledWith(messageId, state);
+      expect(mockMessageModel.query).toHaveBeenCalled();
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+  });
+
+  describe('updateMessage', () => {
+    it('should update message and return { success: true } when no sessionId/topicId provided', async () => {
+      const messageId = 'msg-1';
+      const value = { content: 'updated content' };
+
+      const result = await messageService.updateMessage(messageId, value as any, {});
+
+      expect(mockMessageModel.update).toHaveBeenCalledWith(messageId, value);
+      expect(result).toEqual({ success: true });
+      expect(mockMessageModel.query).not.toHaveBeenCalled();
+    });
+
+    it('should update message and return message list when sessionId provided', async () => {
+      const messageId = 'msg-1';
+      const value = { content: 'updated content' };
+      const mockMessages = [{ id: 'msg-1', content: 'updated content' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.updateMessage(messageId, value as any, {
+        sessionId: 'session-1',
+      });
+
+      expect(mockMessageModel.update).toHaveBeenCalledWith(messageId, value);
+      expect(mockMessageModel.query).toHaveBeenCalled();
+      expect(result).toEqual({ messages: mockMessages, success: true });
+    });
+  });
+
+  describe('useGroup option', () => {
+    it('should pass useGroup option to query', async () => {
+      const mockMessages = [{ id: 'msg-1', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      await messageService.removeMessage('msg-1', {
+        sessionId: 'session-1',
+        useGroup: true,
+      });
+
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.objectContaining({
+          groupAssistantMessages: true,
+        }),
+      );
+    });
+
+    it('should default useGroup to false', async () => {
+      const mockMessages = [{ id: 'msg-1', content: 'test' }];
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      await messageService.removeMessage('msg-1', { sessionId: 'session-1' });
+
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.objectContaining({
+          groupAssistantMessages: false,
+        }),
+      );
+    });
+  });
+
+  describe('createNewMessage', () => {
+    it('should create message and return message list', async () => {
+      const params = {
+        content: 'Hello',
+        role: 'user' as const,
+        sessionId: 'session-1',
+      };
+      const createdMessage = { id: 'msg-1', ...params };
+      const mockMessages = [createdMessage, { id: 'msg-2', content: 'Hi' }];
+
+      vi.mocked(mockMessageModel.create).mockResolvedValue(createdMessage as any);
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.createNewMessage(params as any);
+
+      expect(mockMessageModel.create).toHaveBeenCalledWith(params);
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        {
+          current: 0,
+          groupId: undefined,
+          pageSize: 9999,
+          sessionId: 'session-1',
+          topicId: undefined,
+        },
+        expect.objectContaining({
+          groupAssistantMessages: false,
+        }),
+      );
+      expect(result).toEqual({
+        id: 'msg-1',
+        messages: mockMessages,
+      });
+    });
+
+    it('should create message with useGroup option', async () => {
+      const params = {
+        content: 'Hello',
+        role: 'assistant' as const,
+        sessionId: 'session-1',
+      };
+      const createdMessage = { id: 'msg-1', ...params };
+      const mockMessages = [createdMessage];
+
+      vi.mocked(mockMessageModel.create).mockResolvedValue(createdMessage as any);
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.createNewMessage(params as any, { useGroup: true });
+
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        expect.anything(),
+        expect.objectContaining({
+          groupAssistantMessages: true,
+        }),
+      );
+      expect(result).toEqual({
+        id: 'msg-1',
+        messages: mockMessages,
+      });
+    });
+
+    it('should create message with topicId and groupId', async () => {
+      const params = {
+        content: 'Hello',
+        groupId: 'group-1',
+        role: 'user' as const,
+        sessionId: 'session-1',
+        topicId: 'topic-1',
+      };
+      const createdMessage = { id: 'msg-1', ...params };
+      const mockMessages = [createdMessage];
+
+      vi.mocked(mockMessageModel.create).mockResolvedValue(createdMessage as any);
+      vi.mocked(mockMessageModel.query).mockResolvedValue(mockMessages as any);
+
+      const result = await messageService.createNewMessage(params as any);
+
+      expect(mockMessageModel.query).toHaveBeenCalledWith(
+        {
+          current: 0,
+          groupId: 'group-1',
+          pageSize: 9999,
+          sessionId: 'session-1',
+          topicId: 'topic-1',
+        },
+        expect.anything(),
+      );
+      expect(result.id).toBe('msg-1');
+      expect(result.messages).toEqual(mockMessages);
+    });
+  });
+});

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

@@ -0,0 +1,159 @@
+import { LobeChatDatabase } from '@lobechat/database';
+import { CreateMessageParams, UpdateMessageParams } from '@lobechat/types';
+
+import { MessageModel } from '@/database/models/message';
+
+import { FileService } from '../file';
+
+interface QueryOptions {
+  groupId?: string | null;
+  sessionId?: string | null;
+  topicId?: string | null;
+  useGroup?: boolean;
+}
+
+interface CreateMessageResult {
+  id: string;
+  messages: any[];
+}
+
+/**
+ * Message Service
+ *
+ * Encapsulates repeated "mutation + conditional query" logic.
+ * After performing update/delete operations, conditionally returns message list based on sessionId/topicId.
+ */
+export class MessageService {
+  private messageModel: MessageModel;
+  private fileService: FileService;
+
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.messageModel = new MessageModel(db, userId);
+    this.fileService = new FileService(db, userId);
+  }
+
+  /**
+   * Unified URL processing function
+   */
+  private get postProcessUrl() {
+    return (path: string | null) => this.fileService.getFullFileUrl(path);
+  }
+
+  /**
+   * Unified query options
+   */
+  private getQueryOptions(options: QueryOptions) {
+    return {
+      groupAssistantMessages: options.useGroup ?? false,
+      postProcessUrl: this.postProcessUrl,
+    };
+  }
+
+  /**
+   * Query messages and return response with success status (used after mutations)
+   */
+  private async queryWithSuccess(options?: QueryOptions) {
+    if (!options || (options.sessionId === undefined && options.topicId === undefined)) {
+      return { success: true };
+    }
+
+    const { sessionId, topicId, groupId } = options;
+
+    const messages = await this.messageModel.query(
+      { groupId, sessionId, topicId },
+      this.getQueryOptions(options),
+    );
+
+    return { messages, success: true };
+  }
+
+  /**
+   * Remove messages with optional message list return
+   * Pattern: delete + conditional query
+   */
+  async removeMessages(ids: string[], options?: QueryOptions) {
+    await this.messageModel.deleteMessages(ids);
+    return this.queryWithSuccess(options);
+  }
+
+  /**
+   * Remove single message with optional message list return
+   * Pattern: delete + conditional query
+   */
+  async removeMessage(id: string, options?: QueryOptions) {
+    await this.messageModel.deleteMessage(id);
+    return this.queryWithSuccess(options);
+  }
+
+  /**
+   * Update message RAG with optional message list return
+   * Pattern: update + conditional query
+   */
+  async updateMessageRAG(id: string, value: any, options?: QueryOptions) {
+    await this.messageModel.updateMessageRAG(id, value);
+    return this.queryWithSuccess(options);
+  }
+
+  /**
+   * Update plugin error with optional message list return
+   * Pattern: update + conditional query
+   */
+  async updatePluginError(id: string, value: any, options?: QueryOptions) {
+    await this.messageModel.updateMessagePlugin(id, { error: value });
+    return this.queryWithSuccess(options);
+  }
+
+  /**
+   * Update plugin state and return message list
+   * Pattern: update + conditional query
+   */
+  async updatePluginState(id: string, value: any, options: QueryOptions): Promise<any> {
+    await this.messageModel.updatePluginState(id, value);
+    return this.queryWithSuccess(options);
+  }
+
+  /**
+   * Update message and return message list
+   * Pattern: update + conditional query
+   */
+  async updateMessage(id: string, value: UpdateMessageParams, options: QueryOptions): Promise<any> {
+    await this.messageModel.update(id, value as any);
+    return this.queryWithSuccess(options);
+  }
+
+  /**
+   * Create a new message and return the complete message list
+   * Pattern: create + query
+   *
+   * This method combines message creation and querying into a single operation,
+   * reducing the need for separate refresh calls and improving performance.
+   */
+  async createNewMessage(
+    params: CreateMessageParams,
+    options?: QueryOptions,
+  ): Promise<CreateMessageResult> {
+    // 1. Create the message
+    const item = await this.messageModel.create(params);
+
+    // 2. Query all messages for this session/topic
+    const messages = await this.messageModel.query(
+      {
+        current: 0,
+        groupId: params.groupId,
+        pageSize: 9999,
+        sessionId: params.sessionId,
+        topicId: params.topicId,
+      },
+      {
+        groupAssistantMessages: options?.useGroup ?? false,
+        postProcessUrl: this.postProcessUrl,
+      },
+    );
+
+    // 3. Return the result
+    return {
+      id: item.id,
+      messages,
+    };
+  }
+}

package/src/services/message/index.ts

@@ -31,6 +31,7 @@ export class MessageService {
     return lambdaClient.message.createNewMessage.mutate({
       ...params,
       sessionId: sessionId ? this.toDbSessionId(sessionId) : undefined,
+      useGroup: this.useGroup,
     });
   };