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

package/src/store/chat/slices/aiChat/actions/__tests__/generateAIChatV2.test.ts

@@ -19,8 +19,33 @@ import { resetTestEnvironment, setupMockSelectors, spyOnMessageService } from '.
 // Keep zustand mock as it's needed globally
 vi.mock('zustand/traditional');
 
+// Mock AntdStaticMethods
+vi.mock('@/components/AntdStaticMethods', () => ({
+  notification: {
+    error: vi.fn(),
+    success: vi.fn(),
+    info: vi.fn(),
+    warning: vi.fn(),
+  },
+  message: {
+    error: vi.fn(),
+    success: vi.fn(),
+    info: vi.fn(),
+    warning: vi.fn(),
+  },
+}));
+
+// Mock sessionService to prevent TRPC requests
+vi.mock('@/services/session', () => ({
+  sessionService: {
+    updateSession: vi.fn(),
+    updateSessionConfig: vi.fn(),
+    updateSessionChatConfig: vi.fn(),
+  },
+}));
+
 // Mock server mode for V2 tests
-vi.mock('@/const/version', async (importOriginal) => {
+vi.mock('@lobechat/const', async (importOriginal) => {
   const module = await importOriginal();
   return {
     ...(module as any),
@@ -106,7 +131,7 @@ describe('generateAIChatV2 actions', () => {
           await result.current.sendMessage({ message: TEST_CONTENT.USER_MESSAGE });
         });
 
-        expect(messageService.createMessage).not.toHaveBeenCalled();
+        expect(messageService.createNewMessage).not.toHaveBeenCalled();
         expect(result.current.internal_execAgentRuntime).not.toHaveBeenCalled();
       });
 
@@ -117,7 +142,7 @@ describe('generateAIChatV2 actions', () => {
           await result.current.sendMessage({ message: TEST_CONTENT.EMPTY });
         });
 
-        expect(messageService.createMessage).not.toHaveBeenCalled();
+        expect(messageService.createNewMessage).not.toHaveBeenCalled();
       });
 
       it('should not send when message is empty with empty files array', async () => {
@@ -127,7 +152,7 @@ describe('generateAIChatV2 actions', () => {
           await result.current.sendMessage({ message: TEST_CONTENT.EMPTY, files: [] });
         });
 
-        expect(messageService.createMessage).not.toHaveBeenCalled();
+        expect(messageService.createNewMessage).not.toHaveBeenCalled();
       });
     });
 
@@ -287,7 +312,7 @@ describe('generateAIChatV2 actions', () => {
           });
         });
 
-        expect(messageService.createMessage).toHaveBeenCalled();
+        expect(messageService.createNewMessage).toHaveBeenCalled();
         expect(result.current.internal_execAgentRuntime).not.toHaveBeenCalled();
       });
 
@@ -516,7 +541,7 @@ describe('generateAIChatV2 actions', () => {
         result.current.cancelSendMessageInServer();
       });
 
-      expect(mockAbort).toHaveBeenCalledWith('User cancelled sendMessageInServer operation');
+      expect(mockAbort).toHaveBeenCalledWith('User cancelled sendMessage operation');
       expect(
         result.current.mainSendMessageOperations[
           messageMapKey(TEST_IDS.SESSION_ID, TEST_IDS.TOPIC_ID)
@@ -546,7 +571,7 @@ describe('generateAIChatV2 actions', () => {
         result.current.cancelSendMessageInServer(customTopicId);
       });
 
-      expect(mockAbort).toHaveBeenCalledWith('User cancelled sendMessageInServer operation');
+      expect(mockAbort).toHaveBeenCalledWith('User cancelled sendMessage operation');
     });
 
     it('should handle gracefully when operation does not exist', () => {
@@ -715,4 +740,311 @@ describe('generateAIChatV2 actions', () => {
       });
     });
   });
+
+  describe('callToolFollowAssistantMessage', () => {
+    const TOOL_RESULT_MSG_ID = 'tool-result-msg-id';
+    const ASSISTANT_BLOCK_ID = 'assistant-block-id';
+    const GROUP_MESSAGE_ID = 'group-message-id';
+    const TOOL_CALL_ID = 'tool-call-id';
+
+    beforeEach(() => {
+      // Reset mocks
+      vi.spyOn(messageService, 'createNewMessage').mockResolvedValue({
+        id: 'new-assistant-block-id',
+        messages: [] as any,
+      });
+    });
+
+    it('should find group message from tool result_msg_id in tools array', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const dispatchSpy = vi.fn();
+
+      // Create a group message structure with tool results
+      const groupMessage: UIChatMessage = {
+        id: GROUP_MESSAGE_ID,
+        role: 'group',
+        content: '',
+        sessionId: TEST_IDS.SESSION_ID,
+        topicId: TEST_IDS.TOPIC_ID,
+        children: [
+          {
+            id: ASSISTANT_BLOCK_ID,
+            content: 'Assistant response',
+            tools: [
+              {
+                id: TOOL_CALL_ID,
+                type: 'builtin',
+                apiName: 'testTool',
+                identifier: 'test-tool',
+                arguments: '{}',
+                result: {
+                  id: TOOL_RESULT_MSG_ID,
+                  content: 'Tool result',
+                },
+                result_msg_id: TOOL_RESULT_MSG_ID,
+              },
+            ],
+          },
+        ],
+      } as any;
+
+      act(() => {
+        useChatStore.setState({
+          activeId: TEST_IDS.SESSION_ID,
+          activeTopicId: TEST_IDS.TOPIC_ID,
+          messagesMap: {
+            [messageMapKey(TEST_IDS.SESSION_ID, TEST_IDS.TOPIC_ID)]: [groupMessage],
+          },
+          internal_execAgentRuntime: vi.fn(),
+          internal_dispatchMessage: dispatchSpy,
+        });
+      });
+
+      await act(async () => {
+        await result.current.callToolFollowAssistantMessage({
+          parentId: TOOL_RESULT_MSG_ID,
+        });
+      });
+
+      // Verify that addGroupBlock was called with the correct groupMessageId
+      expect(dispatchSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          type: 'addGroupBlock',
+          groupMessageId: GROUP_MESSAGE_ID,
+        }),
+      );
+
+      // Verify that createNewMessage was called with message params
+      expect(messageService.createNewMessage).toHaveBeenCalledWith(
+        expect.objectContaining({
+          role: 'assistant',
+          parentId: TOOL_RESULT_MSG_ID,
+        }),
+      );
+    });
+
+    it('should handle case when tool result is not found in any group message', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const dispatchSpy = vi.fn();
+
+      const groupMessage: UIChatMessage = {
+        id: GROUP_MESSAGE_ID,
+        role: 'group',
+        content: '',
+        sessionId: TEST_IDS.SESSION_ID,
+        children: [
+          {
+            id: ASSISTANT_BLOCK_ID,
+            content: 'Assistant response',
+            tools: [], // No tools
+          },
+        ],
+      } as any;
+
+      act(() => {
+        useChatStore.setState({
+          activeId: TEST_IDS.SESSION_ID,
+          activeTopicId: TEST_IDS.TOPIC_ID,
+          messagesMap: {
+            [messageMapKey(TEST_IDS.SESSION_ID, TEST_IDS.TOPIC_ID)]: [groupMessage],
+          },
+          internal_execAgentRuntime: vi.fn(),
+          internal_dispatchMessage: dispatchSpy,
+        });
+      });
+
+      await act(async () => {
+        await result.current.callToolFollowAssistantMessage({
+          parentId: 'non-existent-tool-result-id',
+        });
+      });
+
+      // Should create message as regular top-level message (createMessage, not addGroupBlock)
+      expect(dispatchSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          type: 'createMessage',
+        }),
+      );
+
+      expect(messageService.createNewMessage).toHaveBeenCalledWith(
+        expect.objectContaining({
+          role: 'assistant',
+          parentId: 'non-existent-tool-result-id',
+        }),
+      );
+    });
+
+    it('should find group message from nested tool results in multiple children', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const dispatchSpy = vi.fn();
+
+      const groupMessage: UIChatMessage = {
+        id: GROUP_MESSAGE_ID,
+        role: 'group',
+        content: '',
+        sessionId: TEST_IDS.SESSION_ID,
+        children: [
+          {
+            id: 'first-block',
+            content: 'First assistant response',
+            tools: [
+              {
+                id: 'tool-1',
+                type: 'builtin',
+                apiName: 'tool1',
+                identifier: 'tool-1',
+                arguments: '{}',
+                result_msg_id: 'other-result-id',
+              },
+            ],
+          },
+          {
+            id: 'second-block',
+            content: 'Second assistant response',
+            tools: [
+              {
+                id: 'tool-2',
+                type: 'builtin',
+                apiName: 'tool2',
+                identifier: 'tool-2',
+                arguments: '{}',
+                result_msg_id: TOOL_RESULT_MSG_ID, // Target tool result
+              },
+            ],
+          },
+        ],
+      } as any;
+
+      act(() => {
+        useChatStore.setState({
+          activeId: TEST_IDS.SESSION_ID,
+          messagesMap: {
+            [messageMapKey(TEST_IDS.SESSION_ID, TEST_IDS.TOPIC_ID)]: [groupMessage],
+          },
+          internal_execAgentRuntime: vi.fn(),
+          internal_dispatchMessage: dispatchSpy,
+        });
+      });
+
+      await act(async () => {
+        await result.current.callToolFollowAssistantMessage({
+          parentId: TOOL_RESULT_MSG_ID,
+        });
+      });
+
+      // Should find the correct group message even with multiple children
+      expect(dispatchSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          type: 'addGroupBlock',
+          groupMessageId: GROUP_MESSAGE_ID,
+        }),
+      );
+    });
+
+    it('should call internal_execAgentRuntime after creating assistant message', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const mockExecAgentRuntime = vi.fn();
+
+      const groupMessage: UIChatMessage = {
+        id: GROUP_MESSAGE_ID,
+        role: 'group',
+        content: '',
+        sessionId: TEST_IDS.SESSION_ID,
+        children: [
+          {
+            id: ASSISTANT_BLOCK_ID,
+            content: 'Response',
+            tools: [
+              {
+                id: TOOL_CALL_ID,
+                type: 'builtin',
+                apiName: 'test',
+                identifier: 'test',
+                arguments: '{}',
+                result_msg_id: TOOL_RESULT_MSG_ID,
+              },
+            ],
+          },
+        ],
+      } as any;
+
+      act(() => {
+        useChatStore.setState({
+          activeId: TEST_IDS.SESSION_ID,
+          messagesMap: {
+            [messageMapKey(TEST_IDS.SESSION_ID, TEST_IDS.TOPIC_ID)]: [groupMessage],
+          },
+          internal_execAgentRuntime: mockExecAgentRuntime,
+        });
+      });
+
+      await act(async () => {
+        await result.current.callToolFollowAssistantMessage({
+          parentId: TOOL_RESULT_MSG_ID,
+          traceId: 'test-trace-id',
+          threadId: 'test-thread-id',
+        });
+      });
+
+      expect(mockExecAgentRuntime).toHaveBeenCalledWith(
+        expect.objectContaining({
+          assistantMessageId: 'new-assistant-block-id',
+          traceId: 'test-trace-id',
+          threadId: 'test-thread-id',
+        }),
+      );
+    });
+
+    it('should handle missing result_msg_id field gracefully', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const dispatchSpy = vi.fn();
+
+      const groupMessage: UIChatMessage = {
+        id: GROUP_MESSAGE_ID,
+        role: 'group',
+        content: '',
+        sessionId: TEST_IDS.SESSION_ID,
+        children: [
+          {
+            id: ASSISTANT_BLOCK_ID,
+            content: 'Response',
+            tools: [
+              {
+                id: TOOL_CALL_ID,
+                type: 'builtin',
+                apiName: 'test',
+                identifier: 'test',
+                arguments: '{}',
+                // Missing result_msg_id
+              },
+            ],
+          },
+        ],
+      } as any;
+
+      act(() => {
+        useChatStore.setState({
+          activeId: TEST_IDS.SESSION_ID,
+          messagesMap: {
+            [messageMapKey(TEST_IDS.SESSION_ID, TEST_IDS.TOPIC_ID)]: [groupMessage],
+          },
+          internal_execAgentRuntime: vi.fn(),
+          internal_dispatchMessage: dispatchSpy,
+        });
+      });
+
+      await act(async () => {
+        await result.current.callToolFollowAssistantMessage({
+          parentId: TOOL_RESULT_MSG_ID,
+        });
+      });
+
+      // Should create as regular message since no groupMessageId found
+      expect(dispatchSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          type: 'createMessage',
+        }),
+      );
+    });
+  });
 });

package/src/store/chat/slices/aiChat/actions/__tests__/helpers.ts

@@ -58,13 +58,17 @@ export const createMockAbortController = () => {
  */
 export const spyOnMessageService = () => {
   const createMessageSpy = vi
-    .spyOn(messageService, 'createMessage')
-    .mockResolvedValue(TEST_IDS.NEW_MESSAGE_ID);
-  const updateMessageSpy = vi.spyOn(messageService, 'updateMessage').mockResolvedValue(undefined);
-  const removeMessageSpy = vi.spyOn(messageService, 'removeMessage').mockResolvedValue(undefined);
+    .spyOn(messageService, 'createNewMessage')
+    .mockResolvedValue({ id: TEST_IDS.NEW_MESSAGE_ID, messages: [] });
+  const updateMessageSpy = vi
+    .spyOn(messageService, 'updateMessage')
+    .mockResolvedValue({ messages: [], success: true });
+  const removeMessageSpy = vi
+    .spyOn(messageService, 'removeMessage')
+    .mockResolvedValue(undefined as any);
   const updateMessageErrorSpy = vi
     .spyOn(messageService, 'updateMessageError')
-    .mockResolvedValue(undefined);
+    .mockResolvedValue(undefined as any);
 
   return {
     createMessageSpy,