@lobehub/lobehub 2.0.0-next.3 → 2.0.0-next.31

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

@@ -72,7 +72,7 @@ describe('Message Router Integration Tests', () => {
     it('should create message with correct sessionId and topicId', async () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
 
-      const messageId = await caller.createMessage({
+      const result = await caller.createNewMessage({
         content: 'Test message',
         role: 'user',
         sessionId: testSessionId,
@@ -83,11 +83,11 @@ describe('Message Router Integration Tests', () => {
       const [createdMessage] = await serverDB
         .select()
         .from(messages)
-        .where(eq(messages.id, messageId));
+        .where(eq(messages.id, result.id));
 
       expect(createdMessage).toBeDefined();
       expect(createdMessage).toMatchObject({
-        id: messageId,
+        id: result.id,
         sessionId: testSessionId,
         topicId: testTopicId,
         userId: userId,
@@ -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,9 +109,9 @@ describe('Message Router Integration Tests', () => {
           sourceMessageId: 'msg-source',
           type: 'continuation', // type is required
         })
-        .returning();
+        .returning()) as any;
 
-      const messageId = await caller.createMessage({
+      const result = await caller.createNewMessage({
         content: 'Test message in thread',
         role: 'user',
         sessionId: testSessionId,
@@ -123,12 +123,12 @@ describe('Message Router Integration Tests', () => {
       const [createdMessage] = await serverDB
         .select()
         .from(messages)
-        .where(eq(messages.id, messageId));
+        .where(eq(messages.id, result.id));
 
       expect(createdMessage).toBeDefined();
       expect(createdMessage.threadId).toBe(thread.id);
       expect(createdMessage).toMatchObject({
-        id: messageId,
+        id: result.id,
         sessionId: testSessionId,
         topicId: testTopicId,
         threadId: thread.id,
@@ -140,7 +140,7 @@ describe('Message Router Integration Tests', () => {
     it('should create message without topicId', async () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
 
-      const messageId = await caller.createMessage({
+      const result = await caller.createNewMessage({
         content: 'Test message without topic',
         role: 'user',
         sessionId: testSessionId,
@@ -150,7 +150,7 @@ describe('Message Router Integration Tests', () => {
       const [createdMessage] = await serverDB
         .select()
         .from(messages)
-        .where(eq(messages.id, messageId));
+        .where(eq(messages.id, result.id));
 
       expect(createdMessage.topicId).toBeNull();
       expect(createdMessage.sessionId).toBe(testSessionId);
@@ -160,7 +160,7 @@ describe('Message Router Integration Tests', () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
 
       await expect(
-        caller.createMessage({
+        caller.createNewMessage({
           content: 'Test message',
           role: 'user',
           sessionId: 'non-existent-session',
@@ -192,7 +192,7 @@ describe('Message Router Integration Tests', () => {
 
       // 尝试在 testSessionId 下创建消息，但使用 anotherTopic 的 ID
       await expect(
-        caller.createMessage({
+        caller.createNewMessage({
           content: 'Test message',
           role: 'user',
           sessionId: testSessionId,
@@ -207,13 +207,13 @@ describe('Message Router Integration Tests', () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
 
       // 创建多个消息
-      const msg1Id = await caller.createMessage({
+      const msg1Result = await caller.createNewMessage({
         content: 'Message 1',
         role: 'user',
         sessionId: testSessionId,
       });
 
-      const msg2Id = await caller.createMessage({
+      const msg2Result = await caller.createNewMessage({
         content: 'Message 2',
         role: 'assistant',
         sessionId: testSessionId,
@@ -228,7 +228,7 @@ describe('Message Router Integration Tests', () => {
         })
         .returning();
 
-      await caller.createMessage({
+      await caller.createNewMessage({
         content: 'Message in another session',
         role: 'user',
         sessionId: anotherSession.id,
@@ -240,15 +240,15 @@ describe('Message Router Integration Tests', () => {
       });
 
       expect(result).toHaveLength(2);
-      expect(result.map((m) => m.id)).toContain(msg1Id);
-      expect(result.map((m) => m.id)).toContain(msg2Id);
+      expect(result.map((m) => m.id)).toContain(msg1Result.id);
+      expect(result.map((m) => m.id)).toContain(msg2Result.id);
     });
 
     it('should return messages filtered by topicId', async () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
 
       // 在 topic 中创建消息
-      const msgInTopicId = await caller.createMessage({
+      const msgInTopicResult = await caller.createNewMessage({
         content: 'Message in topic',
         role: 'user',
         sessionId: testSessionId,
@@ -256,7 +256,7 @@ describe('Message Router Integration Tests', () => {
       });
 
       // 在 session 中创建消息（不在 topic 中）
-      await caller.createMessage({
+      await caller.createNewMessage({
         content: 'Message without topic',
         role: 'user',
         sessionId: testSessionId,
@@ -269,7 +269,7 @@ describe('Message Router Integration Tests', () => {
       });
 
       expect(result).toHaveLength(1);
-      expect(result[0].id).toBe(msgInTopicId);
+      expect(result[0].id).toBe(msgInTopicResult.id);
       expect(result[0].topicId).toBe(testTopicId);
     });
 
@@ -278,7 +278,7 @@ describe('Message Router Integration Tests', () => {
 
       // 创建多个消息
       for (let i = 0; i < 5; i++) {
-        await caller.createMessage({
+        await caller.createNewMessage({
           content: `Pagination test message ${i}`,
           role: 'user',
           sessionId: testSessionId,
@@ -316,46 +316,73 @@ describe('Message Router Integration Tests', () => {
         expect(page1Ids).not.toEqual(page2Ids);
       }
     });
-  });
 
-  describe('batchCreateMessages', () => {
-    it('should create multiple messages in batch', async () => {
+    it('should return messages filtered by groupId', async () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
 
-      const messagesToCreate = [
-        {
-          content: 'Batch message 1',
-          role: 'user' as const,
-          sessionId: testSessionId,
-        },
-        {
-          content: 'Batch message 2',
-          role: 'assistant' as const,
-          sessionId: testSessionId,
-        },
-        {
-          content: 'Batch message 3',
-          role: 'user' as const,
-          sessionId: testSessionId,
-          topicId: testTopicId,
-        },
-      ];
+      // 首先创建一个 chat_group
+      const { chatGroups } = await import('@/database/schemas');
+      const [chatGroup] = await serverDB
+        .insert(chatGroups)
+        .values({
+          userId,
+          title: 'Test Chat Group',
+        })
+        .returning();
 
-      const result = await caller.batchCreateMessages(messagesToCreate);
+      // 创建消息并设置 groupId
+      const msg1 = await caller.createNewMessage({
+        content: 'Message 1 in group',
+        role: 'assistant',
+        sessionId: testSessionId,
+      });
 
-      expect(result.success).toBe(true);
-      // Note: rowCount might be undefined in PGlite, so we skip this check
-      // expect(result.added).toBe(3);
+      await serverDB
+        .update(messages)
+        .set({ groupId: chatGroup.id })
+        .where(eq(messages.id, msg1.id));
 
-      // 验证数据库中的消息
-      const dbMessages = await serverDB
-        .select()
-        .from(messages)
-        .where(eq(messages.sessionId, testSessionId));
+      // 创建不在 group 中的消息
+      await caller.createNewMessage({
+        content: 'Message without group',
+        role: 'user',
+        sessionId: testSessionId,
+      });
 
-      expect(dbMessages.length).toBeGreaterThanOrEqual(3);
-      const topicMessage = dbMessages.find((m) => m.content === 'Batch message 3');
-      expect(topicMessage?.topicId).toBe(testTopicId);
+      // 查询 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);
     });
   });
 
@@ -364,20 +391,20 @@ describe('Message Router Integration Tests', () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
 
       // 创建消息
-      const msg1Id = await caller.createMessage({
+      const msg1Result = await caller.createNewMessage({
         content: 'Message 1',
         role: 'user',
         sessionId: testSessionId,
       });
 
-      const msg2Id = await caller.createMessage({
+      const msg2Result = await caller.createNewMessage({
         content: 'Message 2',
         role: 'user',
         sessionId: testSessionId,
       });
 
       // 删除消息
-      await caller.removeMessages({ ids: [msg1Id, msg2Id] });
+      await caller.removeMessages({ ids: [msg1Result.id, msg2Result.id] });
 
       // 验证消息已删除
       const remainingMessages = await serverDB
@@ -387,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', () => {
@@ -394,13 +575,13 @@ describe('Message Router Integration Tests', () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
 
       // 创建多个消息
-      await caller.createMessage({
+      await caller.createNewMessage({
         content: 'Message 1',
         role: 'user',
         sessionId: testSessionId,
       });
 
-      await caller.createMessage({
+      await caller.createNewMessage({
         content: 'Message 2',
         role: 'assistant',
         sessionId: testSessionId,
@@ -424,7 +605,7 @@ describe('Message Router Integration Tests', () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
 
       // 在 topic 中创建消息
-      await caller.createMessage({
+      await caller.createNewMessage({
         content: 'Message in topic',
         role: 'user',
         sessionId: testSessionId,
@@ -432,7 +613,7 @@ describe('Message Router Integration Tests', () => {
       });
 
       // 在 session 中创建消息（不在 topic 中）
-      const msgOutsideTopicId = await caller.createMessage({
+      const msgOutsideTopicResult = await caller.createNewMessage({
         content: 'Message outside topic',
         role: 'user',
         sessionId: testSessionId,
@@ -451,7 +632,44 @@ describe('Message Router Integration Tests', () => {
         .where(eq(messages.sessionId, testSessionId));
 
       expect(remainingMessages).toHaveLength(1);
-      expect(remainingMessages[0].id).toBe(msgOutsideTopicId);
+      expect(remainingMessages[0].id).toBe(msgOutsideTopicResult.id);
+    });
+  });
+
+  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();
     });
   });
 
@@ -459,14 +677,14 @@ describe('Message Router Integration Tests', () => {
     it('should update message content', async () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
 
-      const messageId = await caller.createMessage({
+      const result = await caller.createNewMessage({
         content: 'Original content',
         role: 'user',
         sessionId: testSessionId,
       });
 
       await caller.update({
-        id: messageId,
+        id: result.id,
         value: {
           content: 'Updated content',
         },
@@ -475,23 +693,53 @@ describe('Message Router Integration Tests', () => {
       const [updatedMessage] = await serverDB
         .select()
         .from(messages)
-        .where(eq(messages.id, messageId));
+        .where(eq(messages.id, result.id));
 
       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', () => {
     it('should search messages by keyword', async () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
 
-      await caller.createMessage({
+      await caller.createNewMessage({
         content: 'This is a test message about TypeScript',
         role: 'user',
         sessionId: testSessionId,
       });
 
-      await caller.createMessage({
+      await caller.createNewMessage({
         content: 'Another message about JavaScript',
         role: 'user',
         sessionId: testSessionId,
@@ -506,18 +754,470 @@ 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));
 
       // 创建消息
-      await caller.createMessage({
+      await caller.createNewMessage({
         content: 'Message 1',
         role: 'user',
         sessionId: testSessionId,
       });
 
-      await caller.createMessage({
+      await caller.createNewMessage({
         content: 'Message 2',
         role: 'assistant',
         sessionId: testSessionId,
@@ -528,10 +1228,30 @@ 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));
 
-      await caller.createMessage({
+      await caller.createNewMessage({
         content: 'Hello world',
         role: 'user',
         sessionId: testSessionId,
@@ -541,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);
+    });
   });
 });