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

package/src/server/routers/lambda/__tests__/integration/message.integration.test.ts

@@ -101,7 +101,7 @@ describe('Message Router Integration Tests', () => {
 
       // 先创建 thread
       const { threads } = await import('@/database/schemas');
-      const [thread] = await serverDB
+      const [thread] = (await serverDB
         .insert(threads)
         .values({
           userId,
@@ -109,7 +109,7 @@ describe('Message Router Integration Tests', () => {
           sourceMessageId: 'msg-source',
           type: 'continuation', // type is required
         })
-        .returning();
+        .returning()) as any;
 
       const result = await caller.createNewMessage({
         content: 'Test message in thread',
@@ -316,6 +316,74 @@ describe('Message Router Integration Tests', () => {
         expect(page1Ids).not.toEqual(page2Ids);
       }
     });
+
+    it('should return messages filtered by groupId', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      // 首先创建一个 chat_group
+      const { chatGroups } = await import('@/database/schemas');
+      const [chatGroup] = await serverDB
+        .insert(chatGroups)
+        .values({
+          userId,
+          title: 'Test Chat Group',
+        })
+        .returning();
+
+      // 创建消息并设置 groupId
+      const msg1 = await caller.createNewMessage({
+        content: 'Message 1 in group',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      await serverDB
+        .update(messages)
+        .set({ groupId: chatGroup.id })
+        .where(eq(messages.id, msg1.id));
+
+      // 创建不在 group 中的消息
+      await caller.createNewMessage({
+        content: 'Message without group',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      // 查询 group 中的消息
+      const result = await caller.getMessages({
+        sessionId: testSessionId,
+        groupId: chatGroup.id,
+      });
+
+      expect(result).toHaveLength(1);
+      expect(result[0].id).toBe(msg1.id);
+    });
+
+    it('should support useGroup parameter', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      // 创建多个消息
+      await caller.createNewMessage({
+        content: 'Message 1',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      await caller.createNewMessage({
+        content: 'Message 2',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      // useGroup 参数应该影响消息分组展示
+      const result = await caller.getMessages({
+        sessionId: testSessionId,
+        useGroup: true,
+      });
+
+      expect(result).toBeDefined();
+      expect(Array.isArray(result)).toBe(true);
+    });
   });
 
   describe('removeMessages', () => {
@@ -346,6 +414,160 @@ describe('Message Router Integration Tests', () => {
 
       expect(remainingMessages).toHaveLength(0);
     });
+
+    it('should return message list when sessionId is provided', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      // 创建消息
+      const msg1Result = await caller.createNewMessage({
+        content: 'Message 1',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      const msg2Result = await caller.createNewMessage({
+        content: 'Message 2',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      const msg3Result = await caller.createNewMessage({
+        content: 'Message 3',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      // 删除消息并返回列表
+      const result = await caller.removeMessages({
+        ids: [msg1Result.id],
+        sessionId: testSessionId,
+      });
+
+      expect(result.success).toBe(true);
+      expect(result.messages).toBeDefined();
+      expect(result.messages).toHaveLength(2);
+      expect(result.messages?.map((m) => m.id)).toContain(msg2Result.id);
+      expect(result.messages?.map((m) => m.id)).toContain(msg3Result.id);
+    });
+  });
+
+  describe('removeMessage', () => {
+    it('should remove a single message', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msgResult = await caller.createNewMessage({
+        content: 'Message to remove',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      await caller.removeMessage({ id: msgResult.id });
+
+      // 验证消息已删除
+      const deletedMessage = await serverDB
+        .select()
+        .from(messages)
+        .where(eq(messages.id, msgResult.id));
+
+      expect(deletedMessage).toHaveLength(0);
+    });
+
+    it('should return message list when sessionId is provided', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg1Result = await caller.createNewMessage({
+        content: 'Message 1',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      const msg2Result = await caller.createNewMessage({
+        content: 'Message 2',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      const result = await caller.removeMessage({
+        id: msg1Result.id,
+        sessionId: testSessionId,
+      });
+
+      expect(result.success).toBe(true);
+      expect(result.messages).toBeDefined();
+      expect(result.messages).toHaveLength(1);
+      expect(result.messages?.[0].id).toBe(msg2Result.id);
+    });
+  });
+
+  describe('removeAllMessages', () => {
+    it('should remove all messages for the user', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      // 创建多个 session 和消息
+      await caller.createNewMessage({
+        content: 'Message 1',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      const [anotherSession] = await serverDB
+        .insert(sessions)
+        .values({
+          userId,
+          type: 'agent',
+        })
+        .returning();
+
+      await caller.createNewMessage({
+        content: 'Message 2',
+        role: 'user',
+        sessionId: anotherSession.id,
+      });
+
+      // 删除所有消息
+      await caller.removeAllMessages();
+
+      // 验证所有消息已删除
+      const remainingMessages = await serverDB
+        .select()
+        .from(messages)
+        .where(eq(messages.userId, userId));
+
+      expect(remainingMessages).toHaveLength(0);
+    });
+  });
+
+  describe('removeMessageQuery', () => {
+    it('should remove message query', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msgResult = await caller.createNewMessage({
+        content: 'Message with query',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      // 创建一个 message query 记录，使用 UUID
+      const { messageQueries } = await import('@/database/schemas');
+      const [queryRecord] = await serverDB
+        .insert(messageQueries)
+        .values({
+          messageId: msgResult.id,
+          userId,
+          userQuery: 'test query',
+        })
+        .returning();
+
+      await caller.removeMessageQuery({ id: queryRecord.id });
+
+      // 验证消息查询已删除
+      const deletedQuery = await serverDB
+        .select()
+        .from(messageQueries)
+        .where(eq(messageQueries.id, queryRecord.id));
+
+      expect(deletedQuery).toHaveLength(0);
+    });
   });
 
   describe('removeMessagesByAssistant', () => {
@@ -414,6 +636,43 @@ describe('Message Router Integration Tests', () => {
     });
   });
 
+  describe('removeMessagesByGroup', () => {
+    it('should call removeMessagesByGroup endpoint', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      // 首先创建一个 chat_group
+      const { chatGroups } = await import('@/database/schemas');
+      const [chatGroup] = await serverDB
+        .insert(chatGroups)
+        .values({
+          userId,
+          title: 'Test Chat Group for Delete',
+        })
+        .returning();
+
+      // 创建消息并设置 groupId
+      const msg1 = await caller.createNewMessage({
+        content: 'Message 1 in group',
+        role: 'assistant',
+        sessionId: testSessionId,
+        topicId: testTopicId,
+      });
+
+      await serverDB
+        .update(messages)
+        .set({ groupId: chatGroup.id })
+        .where(eq(messages.id, msg1.id));
+
+      // 调用删除接口（不会抛出错误即为成功）
+      await expect(
+        caller.removeMessagesByGroup({
+          groupId: chatGroup.id,
+          topicId: testTopicId,
+        }),
+      ).resolves.not.toThrow();
+    });
+  });
+
   describe('update', () => {
     it('should update message content', async () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
@@ -438,6 +697,36 @@ describe('Message Router Integration Tests', () => {
 
       expect(updatedMessage.content).toBe('Updated content');
     });
+
+    it('should update message and return message list when sessionId is provided', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg1 = await caller.createNewMessage({
+        content: 'Message 1',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      const msg2 = await caller.createNewMessage({
+        content: 'Message 2',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      const result = await caller.update({
+        id: msg1.id,
+        sessionId: testSessionId,
+        value: {
+          content: 'Updated Message 1',
+        },
+      });
+
+      expect(result).toBeDefined();
+      // The update method returns the updated message list
+      const messages = await caller.getMessages({ sessionId: testSessionId });
+      expect(messages).toHaveLength(2);
+      expect(messages.find((m) => m.id === msg1.id)?.content).toBe('Updated Message 1');
+    });
   });
 
   describe('searchMessages', () => {
@@ -465,6 +754,458 @@ describe('Message Router Integration Tests', () => {
     });
   });
 
+  describe('updateMessagePlugin', () => {
+    it('should update message plugin state', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg = await caller.createNewMessage({
+        content: 'Message with plugin',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      // 先创建一个 plugin 记录
+      const { messagePlugins } = await import('@/database/schemas');
+      await serverDB.insert(messagePlugins).values({
+        id: msg.id,
+        userId,
+        toolCallId: 'test-tool-call',
+        type: 'default',
+      });
+
+      await caller.updateMessagePlugin({
+        id: msg.id,
+        value: {
+          state: { key: 'value' },
+        },
+      });
+
+      const [updatedPlugin] = await serverDB
+        .select()
+        .from(messagePlugins)
+        .where(eq(messagePlugins.id, msg.id));
+
+      expect(updatedPlugin).toBeDefined();
+      expect(updatedPlugin.state).toBeDefined();
+    });
+  });
+
+  describe('updateMessageRAG', () => {
+    it('should update message RAG information', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg = await caller.createNewMessage({
+        content: 'Message with RAG',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      // 创建必要的依赖: chunks -> messageQueries -> messageQueryChunks
+      const { chunks, messageQueries, messageQueryChunks } = await import('@/database/schemas');
+
+      // 1. 创建 chunk
+      const [chunk] = await serverDB
+        .insert(chunks)
+        .values({
+          userId,
+          text: 'test chunk content',
+        })
+        .returning();
+
+      // 2. 创建 message query
+      const [query] = await serverDB
+        .insert(messageQueries)
+        .values({
+          messageId: msg.id,
+          userId,
+          userQuery: 'test query',
+        })
+        .returning();
+
+      // 3. 调用 updateMessageRAG
+      await caller.updateMessageRAG({
+        id: msg.id,
+        value: {
+          fileChunks: [{ id: chunk.id, similarity: 0.95 }],
+          ragQueryId: query.id,
+        },
+      });
+
+      // 验证 messageQueryChunks 记录已创建
+      const [queryChunk] = await serverDB
+        .select()
+        .from(messageQueryChunks)
+        .where(eq(messageQueryChunks.messageId, msg.id));
+
+      expect(queryChunk).toBeDefined();
+      expect(queryChunk.chunkId).toBe(chunk.id);
+    });
+
+    it('should return message list when sessionId is provided', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg1 = await caller.createNewMessage({
+        content: 'Message 1',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      await caller.createNewMessage({
+        content: 'Message 2',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      // 创建必要的依赖: chunks -> messageQueries
+      const { chunks, messageQueries } = await import('@/database/schemas');
+      const [chunk] = await serverDB
+        .insert(chunks)
+        .values({
+          userId,
+          text: 'test chunk content',
+        })
+        .returning();
+
+      // 创建 query (需要 queryId)
+      const [query] = await serverDB
+        .insert(messageQueries)
+        .values({
+          messageId: msg1.id,
+          userId,
+          userQuery: 'test query',
+        })
+        .returning();
+
+      const result = await caller.updateMessageRAG({
+        id: msg1.id,
+        sessionId: testSessionId,
+        value: {
+          fileChunks: [{ id: chunk.id, similarity: 0.95 }],
+          ragQueryId: query.id,
+        },
+      });
+
+      expect(result.success).toBe(true);
+      expect(result.messages).toBeDefined();
+      expect(result.messages).toHaveLength(2);
+    });
+  });
+
+  describe('updateMetadata', () => {
+    it('should update message metadata', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg = await caller.createNewMessage({
+        content: 'Message with metadata',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      await caller.updateMetadata({
+        id: msg.id,
+        value: { customKey: 'customValue' },
+      });
+
+      const [updatedMessage] = await serverDB
+        .select()
+        .from(messages)
+        .where(eq(messages.id, msg.id));
+
+      expect(updatedMessage).toBeDefined();
+      // Verify the message still exists after update
+      expect(updatedMessage.id).toBe(msg.id);
+    });
+  });
+
+  describe('updatePluginError', () => {
+    it('should update plugin error state', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg = await caller.createNewMessage({
+        content: 'Message with plugin error',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      // 先创建一个 plugin 记录
+      const { messagePlugins } = await import('@/database/schemas');
+      await serverDB.insert(messagePlugins).values({
+        id: msg.id,
+        userId,
+        toolCallId: 'test-tool-call-error',
+        type: 'default',
+      });
+
+      await caller.updatePluginError({
+        id: msg.id,
+        value: { message: 'Plugin error occurred' },
+      });
+
+      const [updatedPlugin] = await serverDB
+        .select()
+        .from(messagePlugins)
+        .where(eq(messagePlugins.id, msg.id));
+
+      expect(updatedPlugin).toBeDefined();
+      expect(updatedPlugin.error).toBeDefined();
+    });
+
+    it('should return message list when sessionId is provided', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg1 = await caller.createNewMessage({
+        content: 'Message 1',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      // 先创建一个 plugin 记录
+      const { messagePlugins } = await import('@/database/schemas');
+      await serverDB.insert(messagePlugins).values({
+        id: msg1.id,
+        userId,
+        toolCallId: 'test-tool-call-error-2',
+        type: 'default',
+      });
+
+      await caller.createNewMessage({
+        content: 'Message 2',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      const result = await caller.updatePluginError({
+        id: msg1.id,
+        sessionId: testSessionId,
+        value: { message: 'Error' },
+      });
+
+      expect(result.success).toBe(true);
+      expect(result.messages).toBeDefined();
+      expect(result.messages).toHaveLength(2);
+    });
+  });
+
+  describe('updatePluginState', () => {
+    it('should update plugin state', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg = await caller.createNewMessage({
+        content: 'Message with plugin state',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      // 先创建一个 plugin 记录
+      const { messagePlugins } = await import('@/database/schemas');
+      await serverDB.insert(messagePlugins).values({
+        id: msg.id,
+        userId,
+        toolCallId: 'test-tool-call-state',
+        type: 'default',
+      });
+
+      const result = await caller.updatePluginState({
+        id: msg.id,
+        sessionId: testSessionId,
+        value: { stateKey: 'stateValue' },
+      });
+
+      expect(result).toBeDefined();
+    });
+  });
+
+  describe('updateTTS', () => {
+    it('should update TTS information', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg = await caller.createNewMessage({
+        content: 'Message with TTS',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      // 创建 file 记录
+      const { files } = await import('@/database/schemas');
+      const [file] = await serverDB
+        .insert(files)
+        .values({
+          userId,
+          name: 'audio.mp3',
+          fileType: 'audio/mpeg',
+          size: 1024,
+          url: '/files/audio.mp3',
+        })
+        .returning();
+
+      await caller.updateTTS({
+        id: msg.id,
+        value: {
+          file: file.id,
+          voice: 'en-US-neural',
+          contentMd5: 'abc123',
+        },
+      });
+
+      const { messageTTS } = await import('@/database/schemas');
+      const [ttsRecord] = await serverDB.select().from(messageTTS).where(eq(messageTTS.id, msg.id));
+
+      expect(ttsRecord).toBeDefined();
+      expect(ttsRecord.voice).toBe('en-US-neural');
+      expect(ttsRecord.fileId).toBe(file.id);
+    });
+
+    it('should delete TTS when value is false', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg = await caller.createNewMessage({
+        content: 'Message with TTS to delete',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      // 创建 file 记录
+      const { files } = await import('@/database/schemas');
+      const [file] = await serverDB
+        .insert(files)
+        .values({
+          userId,
+          name: 'audio-delete.mp3',
+          fileType: 'audio/mpeg',
+          size: 1024,
+          url: '/files/audio-delete.mp3',
+        })
+        .returning();
+
+      // First add TTS
+      await caller.updateTTS({
+        id: msg.id,
+        value: {
+          file: file.id,
+          voice: 'en-US-neural',
+        },
+      });
+
+      // Then delete it
+      await caller.updateTTS({
+        id: msg.id,
+        value: false,
+      });
+
+      const { messageTTS } = await import('@/database/schemas');
+      const [ttsRecord] = await serverDB.select().from(messageTTS).where(eq(messageTTS.id, msg.id));
+
+      expect(ttsRecord).toBeUndefined();
+    });
+  });
+
+  describe('updateTranslate', () => {
+    it('should update translation information', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg = await caller.createNewMessage({
+        content: 'Hello world',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      await caller.updateTranslate({
+        id: msg.id,
+        value: {
+          content: '你好世界',
+          from: 'en',
+          to: 'zh',
+        },
+      });
+
+      const { messageTranslates } = await import('@/database/schemas');
+      const [translateRecord] = await serverDB
+        .select()
+        .from(messageTranslates)
+        .where(eq(messageTranslates.id, msg.id));
+
+      expect(translateRecord).toBeDefined();
+      expect(translateRecord.to).toBe('zh');
+    });
+
+    it('should delete translation when value is false', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      const msg = await caller.createNewMessage({
+        content: 'Hello world',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      // First add translation
+      await caller.updateTranslate({
+        id: msg.id,
+        value: {
+          content: '你好世界',
+          to: 'zh',
+        },
+      });
+
+      // Then delete it
+      await caller.updateTranslate({
+        id: msg.id,
+        value: false,
+      });
+
+      const [updatedMessage] = await serverDB
+        .select()
+        .from(messages)
+        .where(eq(messages.id, msg.id));
+
+      expect(updatedMessage).toBeDefined();
+    });
+  });
+
+  describe('getHeatmaps', () => {
+    it('should get message heatmaps', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      // 创建一些消息
+      await caller.createNewMessage({
+        content: 'Message 1',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      await caller.createNewMessage({
+        content: 'Message 2',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      const heatmaps = await caller.getHeatmaps();
+
+      expect(heatmaps).toBeDefined();
+      expect(Array.isArray(heatmaps)).toBe(true);
+    });
+  });
+
+  describe('rankModels', () => {
+    it('should get model usage ranking', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      // 创建带有模型信息的消息
+      const msg = await caller.createNewMessage({
+        content: 'Message from AI',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
+
+      // 添加模型信息
+      await serverDB.update(messages).set({ model: 'gpt-4' }).where(eq(messages.id, msg.id));
+
+      const ranking = await caller.rankModels();
+
+      expect(ranking).toBeDefined();
+      expect(Array.isArray(ranking)).toBe(true);
+    });
+  });
+
   describe('count and statistics', () => {
     it('should count messages', async () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
@@ -487,6 +1228,26 @@ describe('Message Router Integration Tests', () => {
       expect(count).toBe(2);
     });
 
+    it('should count messages with date range', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      await caller.createNewMessage({
+        content: 'Message 1',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      const startDate = new Date(Date.now() - 7 * 24 * 60 * 60 * 1000).toISOString();
+      const endDate = new Date().toISOString();
+
+      const count = await caller.count({
+        startDate,
+        endDate,
+      });
+
+      expect(count).toBeGreaterThanOrEqual(1);
+    });
+
     it('should count words', async () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
 
@@ -500,5 +1261,25 @@ describe('Message Router Integration Tests', () => {
 
       expect(wordCount).toBeGreaterThan(0);
     });
+
+    it('should count words with date range', async () => {
+      const caller = messageRouter.createCaller(createTestContext(userId));
+
+      await caller.createNewMessage({
+        content: 'Hello world test message',
+        role: 'user',
+        sessionId: testSessionId,
+      });
+
+      const startDate = new Date(Date.now() - 7 * 24 * 60 * 60 * 1000).toISOString();
+      const endDate = new Date().toISOString();
+
+      const wordCount = await caller.countWords({
+        startDate,
+        endDate,
+      });
+
+      expect(wordCount).toBeGreaterThan(0);
+    });
   });
 });