@lobehub/lobehub 2.0.0-next.34 → 2.0.0-next.36

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

@@ -0,0 +1,332 @@
+import { act, renderHook } from '@testing-library/react';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { useChatStore } from '../../../../store';
+import { messageMapKey } from '../../../../utils/messageMapKey';
+import { TEST_IDS } from './fixtures';
+import { resetTestEnvironment } from './helpers';
+
+// Keep zustand mock as it's needed globally
+vi.mock('zustand/traditional');
+
+beforeEach(() => {
+  resetTestEnvironment();
+});
+
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+describe('ConversationControl actions', () => {
+  describe('stopGenerateMessage', () => {
+    it('should abort generation and clear loading state when controller exists', () => {
+      const abortController = new AbortController();
+
+      act(() => {
+        useChatStore.setState({ chatLoadingIdsAbortController: abortController });
+      });
+
+      const { result } = renderHook(() => useChatStore());
+      const toggleLoadingSpy = vi.spyOn(result.current, 'internal_toggleChatLoading');
+
+      act(() => {
+        result.current.stopGenerateMessage();
+      });
+
+      expect(abortController.signal.aborted).toBe(true);
+      expect(toggleLoadingSpy).toHaveBeenCalledWith(false, undefined, expect.any(String));
+    });
+
+    it('should do nothing when abort controller is not set', () => {
+      act(() => {
+        useChatStore.setState({ chatLoadingIdsAbortController: undefined });
+      });
+
+      const { result } = renderHook(() => useChatStore());
+      const toggleLoadingSpy = vi.spyOn(result.current, 'internal_toggleChatLoading');
+
+      act(() => {
+        result.current.stopGenerateMessage();
+      });
+
+      expect(toggleLoadingSpy).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('cancelSendMessageInServer', () => {
+    it('should abort operation and restore editor state when cancelling', () => {
+      const { result } = renderHook(() => useChatStore());
+      const mockAbort = vi.fn();
+      const mockSetJSONState = vi.fn();
+
+      act(() => {
+        useChatStore.setState({
+          activeId: TEST_IDS.SESSION_ID,
+          activeTopicId: TEST_IDS.TOPIC_ID,
+          mainSendMessageOperations: {
+            [messageMapKey(TEST_IDS.SESSION_ID, TEST_IDS.TOPIC_ID)]: {
+              isLoading: true,
+              abortController: { abort: mockAbort, signal: {} as any },
+              inputEditorTempState: { content: 'saved content' },
+            },
+          },
+          mainInputEditor: { setJSONState: mockSetJSONState } as any,
+        });
+      });
+
+      act(() => {
+        result.current.cancelSendMessageInServer();
+      });
+
+      expect(mockAbort).toHaveBeenCalledWith('User cancelled sendMessage operation');
+      expect(
+        result.current.mainSendMessageOperations[
+          messageMapKey(TEST_IDS.SESSION_ID, TEST_IDS.TOPIC_ID)
+        ]?.isLoading,
+      ).toBe(false);
+      expect(mockSetJSONState).toHaveBeenCalledWith({ content: 'saved content' });
+    });
+
+    it('should cancel operation for specified topic ID', () => {
+      const { result } = renderHook(() => useChatStore());
+      const mockAbort = vi.fn();
+      const customTopicId = 'custom-topic-id';
+
+      act(() => {
+        useChatStore.setState({
+          activeId: TEST_IDS.SESSION_ID,
+          mainSendMessageOperations: {
+            [messageMapKey(TEST_IDS.SESSION_ID, customTopicId)]: {
+              isLoading: true,
+              abortController: { abort: mockAbort, signal: {} as any },
+            },
+          },
+        });
+      });
+
+      act(() => {
+        result.current.cancelSendMessageInServer(customTopicId);
+      });
+
+      expect(mockAbort).toHaveBeenCalledWith('User cancelled sendMessage operation');
+    });
+
+    it('should handle gracefully when operation does not exist', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({ mainSendMessageOperations: {} });
+      });
+
+      expect(() => {
+        act(() => {
+          result.current.cancelSendMessageInServer('non-existing-topic');
+        });
+      }).not.toThrow();
+    });
+  });
+
+  describe('clearSendMessageError', () => {
+    it('should clear error state for current topic', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({
+          activeId: TEST_IDS.SESSION_ID,
+          activeTopicId: TEST_IDS.TOPIC_ID,
+          mainSendMessageOperations: {
+            [messageMapKey(TEST_IDS.SESSION_ID, TEST_IDS.TOPIC_ID)]: {
+              isLoading: false,
+              inputSendErrorMsg: 'Some error',
+            },
+          },
+        });
+      });
+
+      act(() => {
+        result.current.clearSendMessageError();
+      });
+
+      expect(
+        result.current.mainSendMessageOperations[
+          messageMapKey(TEST_IDS.SESSION_ID, TEST_IDS.TOPIC_ID)
+        ],
+      ).toBeUndefined();
+    });
+
+    it('should handle gracefully when no error operation exists', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({ mainSendMessageOperations: {} });
+      });
+
+      expect(() => {
+        act(() => {
+          result.current.clearSendMessageError();
+        });
+      }).not.toThrow();
+    });
+  });
+
+  describe('internal_toggleSendMessageOperation', () => {
+    it('should create new send operation with abort controller', () => {
+      const { result } = renderHook(() => useChatStore());
+      let abortController: AbortController | undefined;
+
+      act(() => {
+        abortController = result.current.internal_toggleSendMessageOperation('test-key', true);
+      });
+
+      expect(abortController!).toBeInstanceOf(AbortController);
+      expect(result.current.mainSendMessageOperations['test-key']?.isLoading).toBe(true);
+      expect(result.current.mainSendMessageOperations['test-key']?.abortController).toBe(
+        abortController,
+      );
+    });
+
+    it('should stop send operation and clear abort controller', () => {
+      const { result } = renderHook(() => useChatStore());
+      const mockAbortController = { abort: vi.fn() } as any;
+
+      let abortController: AbortController | undefined;
+      act(() => {
+        result.current.internal_updateSendMessageOperation('test-key', {
+          isLoading: true,
+          abortController: mockAbortController,
+        });
+
+        abortController = result.current.internal_toggleSendMessageOperation('test-key', false);
+      });
+
+      expect(abortController).toBeUndefined();
+      expect(result.current.mainSendMessageOperations['test-key']?.isLoading).toBe(false);
+      expect(result.current.mainSendMessageOperations['test-key']?.abortController).toBeNull();
+    });
+
+    it('should call abort with cancel reason when stopping', () => {
+      const { result } = renderHook(() => useChatStore());
+      const mockAbortController = { abort: vi.fn() } as any;
+
+      act(() => {
+        result.current.internal_updateSendMessageOperation('test-key', {
+          isLoading: true,
+          abortController: mockAbortController,
+        });
+
+        result.current.internal_toggleSendMessageOperation('test-key', false, 'Test cancel reason');
+      });
+
+      expect(mockAbortController.abort).toHaveBeenCalledWith('Test cancel reason');
+    });
+
+    it('should support multiple parallel operations', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let abortController1, abortController2;
+      act(() => {
+        abortController1 = result.current.internal_toggleSendMessageOperation('key1', true);
+        abortController2 = result.current.internal_toggleSendMessageOperation('key2', true);
+      });
+
+      expect(result.current.mainSendMessageOperations['key1']?.isLoading).toBe(true);
+      expect(result.current.mainSendMessageOperations['key2']?.isLoading).toBe(true);
+      expect(abortController1).not.toBe(abortController2);
+    });
+  });
+
+  describe('internal_updateSendMessageOperation', () => {
+    it('should update operation state', () => {
+      const { result } = renderHook(() => useChatStore());
+      const mockAbortController = new AbortController();
+
+      act(() => {
+        result.current.internal_updateSendMessageOperation('test-key', {
+          isLoading: true,
+          abortController: mockAbortController,
+          inputSendErrorMsg: 'test error',
+        });
+      });
+
+      expect(result.current.mainSendMessageOperations['test-key']).toEqual({
+        isLoading: true,
+        abortController: mockAbortController,
+        inputSendErrorMsg: 'test error',
+      });
+    });
+
+    it('should support partial update of operation state', () => {
+      const { result } = renderHook(() => useChatStore());
+      const initialController = new AbortController();
+
+      act(() => {
+        result.current.internal_updateSendMessageOperation('test-key', {
+          isLoading: true,
+          abortController: initialController,
+        });
+
+        result.current.internal_updateSendMessageOperation('test-key', {
+          inputSendErrorMsg: 'new error',
+        });
+      });
+
+      expect(result.current.mainSendMessageOperations['test-key']).toEqual({
+        isLoading: true,
+        abortController: initialController,
+        inputSendErrorMsg: 'new error',
+      });
+    });
+  });
+
+  describe('switchMessageBranch', () => {
+    it('should switch to a different message branch', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = TEST_IDS.MESSAGE_ID;
+      const branchIndex = 1;
+
+      const optimisticUpdateSpy = vi
+        .spyOn(result.current, 'optimisticUpdateMessageMetadata')
+        .mockResolvedValue(undefined);
+
+      await act(async () => {
+        await result.current.switchMessageBranch(messageId, branchIndex);
+      });
+
+      expect(optimisticUpdateSpy).toHaveBeenCalledWith(messageId, { activeBranchIndex: branchIndex });
+    });
+
+    it('should handle switching to branch 0', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = TEST_IDS.MESSAGE_ID;
+      const branchIndex = 0;
+
+      const optimisticUpdateSpy = vi
+        .spyOn(result.current, 'optimisticUpdateMessageMetadata')
+        .mockResolvedValue(undefined);
+
+      await act(async () => {
+        await result.current.switchMessageBranch(messageId, branchIndex);
+      });
+
+      expect(optimisticUpdateSpy).toHaveBeenCalledWith(messageId, { activeBranchIndex: 0 });
+    });
+
+    it('should handle errors gracefully when optimistic update fails', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = TEST_IDS.MESSAGE_ID;
+      const branchIndex = 2;
+
+      const optimisticUpdateSpy = vi
+        .spyOn(result.current, 'optimisticUpdateMessageMetadata')
+        .mockRejectedValue(new Error('Update failed'));
+
+      await expect(
+        act(async () => {
+          await result.current.switchMessageBranch(messageId, branchIndex);
+        }),
+      ).rejects.toThrow('Update failed');
+
+      expect(optimisticUpdateSpy).toHaveBeenCalledWith(messageId, { activeBranchIndex: branchIndex });
+    });
+  });
+});

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

@@ -0,0 +1,257 @@
+import { act, renderHook } from '@testing-library/react';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { aiChatService } from '@/services/aiChat';
+
+import { useChatStore } from '../../../../store';
+import { TEST_CONTENT, TEST_IDS, createMockMessage } from './fixtures';
+import {
+  resetTestEnvironment,
+  setupMockSelectors,
+  setupStoreWithMessages,
+  spyOnMessageService,
+} from './helpers';
+
+// Keep zustand mock as it's needed globally
+vi.mock('zustand/traditional');
+
+beforeEach(() => {
+  resetTestEnvironment();
+  setupMockSelectors();
+  spyOnMessageService();
+
+  act(() => {
+    useChatStore.setState({
+      refreshMessages: vi.fn(),
+      refreshTopic: vi.fn(),
+      internal_execAgentRuntime: vi.fn(),
+    });
+  });
+});
+
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+describe('ConversationLifecycle actions', () => {
+  describe('sendMessage', () => {
+    describe('validation', () => {
+      it('should not send when there is no active session', async () => {
+        act(() => {
+          useChatStore.setState({ activeId: undefined });
+        });
+
+        const { result } = renderHook(() => useChatStore());
+
+        await act(async () => {
+          await result.current.sendMessage({ message: TEST_CONTENT.USER_MESSAGE });
+        });
+
+        expect(result.current.internal_execAgentRuntime).not.toHaveBeenCalled();
+      });
+
+      it('should not send when message is empty and no files are provided', async () => {
+        const { result } = renderHook(() => useChatStore());
+
+        await act(async () => {
+          await result.current.sendMessage({ message: TEST_CONTENT.EMPTY });
+        });
+
+        expect(result.current.internal_execAgentRuntime).not.toHaveBeenCalled();
+      });
+
+      it('should not send when message is empty with empty files array', async () => {
+        const { result } = renderHook(() => useChatStore());
+
+        await act(async () => {
+          await result.current.sendMessage({ message: TEST_CONTENT.EMPTY, files: [] });
+        });
+
+        expect(result.current.internal_execAgentRuntime).not.toHaveBeenCalled();
+      });
+    });
+
+    describe('message creation', () => {
+      it('should not process AI when onlyAddUserMessage is true', async () => {
+        const { result } = renderHook(() => useChatStore());
+
+        vi.spyOn(aiChatService, 'sendMessageInServer').mockResolvedValue({
+          messages: [],
+          topics: [],
+          assistantMessageId: TEST_IDS.ASSISTANT_MESSAGE_ID,
+        } as any);
+
+        await act(async () => {
+          await result.current.sendMessage({
+            message: TEST_CONTENT.USER_MESSAGE,
+            onlyAddUserMessage: true,
+          });
+        });
+
+        expect(result.current.internal_execAgentRuntime).not.toHaveBeenCalled();
+      });
+
+      it('should create user message and trigger AI processing', async () => {
+        const { result } = renderHook(() => useChatStore());
+
+        vi.spyOn(aiChatService, 'sendMessageInServer').mockResolvedValue({
+          messages: [
+            createMockMessage({ id: TEST_IDS.USER_MESSAGE_ID, role: 'user' }),
+            createMockMessage({ id: TEST_IDS.ASSISTANT_MESSAGE_ID, role: 'assistant' }),
+          ],
+          topics: [],
+          assistantMessageId: TEST_IDS.ASSISTANT_MESSAGE_ID,
+        } as any);
+
+        await act(async () => {
+          await result.current.sendMessage({ message: TEST_CONTENT.USER_MESSAGE });
+        });
+
+        expect(result.current.internal_execAgentRuntime).toHaveBeenCalled();
+      });
+    });
+  });
+
+  describe('regenerateUserMessage', () => {
+    it('should trigger user message regeneration', async () => {
+      const messages = [
+        createMockMessage({ id: TEST_IDS.USER_MESSAGE_ID, role: 'user', content: 'test' }),
+        createMockMessage({ id: TEST_IDS.MESSAGE_ID, role: 'assistant' }),
+      ];
+
+      setupStoreWithMessages(messages);
+
+      const switchMessageBranchSpy = vi.fn().mockResolvedValue(undefined);
+      const internalTraceSpy = vi.fn();
+
+      act(() => {
+        useChatStore.setState({
+          internal_traceMessage: internalTraceSpy,
+          switchMessageBranch: switchMessageBranchSpy,
+          internal_shouldUseRAG: vi.fn().mockReturnValue(false),
+        });
+      });
+
+      const { result } = renderHook(() => useChatStore());
+
+      await act(async () => {
+        await result.current.regenerateUserMessage(TEST_IDS.USER_MESSAGE_ID);
+      });
+
+      expect(switchMessageBranchSpy).toHaveBeenCalledWith(TEST_IDS.USER_MESSAGE_ID, 1);
+      expect(result.current.internal_execAgentRuntime).toHaveBeenCalledWith(
+        expect.objectContaining({
+          parentMessageId: TEST_IDS.USER_MESSAGE_ID,
+          parentMessageType: 'user',
+        }),
+      );
+      expect(internalTraceSpy).toHaveBeenCalled();
+    });
+
+    it('should not regenerate when already regenerating', async () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({
+          regeneratingIds: [TEST_IDS.USER_MESSAGE_ID],
+          internal_execAgentRuntime: vi.fn(),
+        });
+      });
+
+      await act(async () => {
+        await result.current.regenerateUserMessage(TEST_IDS.USER_MESSAGE_ID);
+      });
+
+      expect(result.current.internal_execAgentRuntime).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('regenerateAssistantMessage', () => {
+    it('should trigger assistant message regeneration', async () => {
+      const messages = [
+        createMockMessage({ id: TEST_IDS.USER_MESSAGE_ID, role: 'user' }),
+        createMockMessage({
+          id: TEST_IDS.MESSAGE_ID,
+          role: 'assistant',
+          parentId: TEST_IDS.USER_MESSAGE_ID,
+        }),
+      ];
+
+      setupStoreWithMessages(messages);
+
+      act(() => {
+        useChatStore.setState({
+          internal_traceMessage: vi.fn(),
+          switchMessageBranch: vi.fn(),
+        });
+      });
+
+      const { result } = renderHook(() => useChatStore());
+
+      await act(async () => {
+        await result.current.regenerateAssistantMessage(TEST_IDS.MESSAGE_ID);
+      });
+
+      expect(result.current.internal_execAgentRuntime).toHaveBeenCalledWith(
+        expect.objectContaining({
+          parentMessageId: TEST_IDS.USER_MESSAGE_ID,
+          parentMessageType: 'user',
+        }),
+      );
+      expect(result.current.internal_traceMessage).toHaveBeenCalled();
+    });
+
+    it('should not regenerate when already regenerating', async () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({
+          regeneratingIds: [TEST_IDS.MESSAGE_ID],
+          internal_execAgentRuntime: vi.fn(),
+        });
+      });
+
+      await act(async () => {
+        await result.current.regenerateAssistantMessage(TEST_IDS.MESSAGE_ID);
+      });
+
+      expect(result.current.internal_execAgentRuntime).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('delAndRegenerateMessage', () => {
+    it('should delete message then regenerate', async () => {
+      const messages = [
+        createMockMessage({ id: TEST_IDS.USER_MESSAGE_ID, role: 'user' }),
+        createMockMessage({
+          id: TEST_IDS.MESSAGE_ID,
+          role: 'assistant',
+          parentId: TEST_IDS.USER_MESSAGE_ID,
+        }),
+      ];
+
+      setupStoreWithMessages(messages);
+
+      act(() => {
+        useChatStore.setState({
+          regenerateAssistantMessage: vi.fn(),
+          deleteMessage: vi.fn(),
+          internal_traceMessage: vi.fn(),
+        });
+      });
+
+      const { result } = renderHook(() => useChatStore());
+
+      await act(async () => {
+        await result.current.delAndRegenerateMessage(TEST_IDS.MESSAGE_ID);
+      });
+
+      expect(result.current.regenerateAssistantMessage).toHaveBeenCalledWith(
+        TEST_IDS.MESSAGE_ID,
+        expect.objectContaining({ skipTrace: true }),
+      );
+      expect(result.current.deleteMessage).toHaveBeenCalledWith(TEST_IDS.MESSAGE_ID);
+      expect(result.current.internal_traceMessage).toHaveBeenCalled();
+    });
+  });
+});

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

@@ -35,11 +35,16 @@ export const setupMockSelectors = (
 /**
  * Setup store state with messages
  */
-export const setupStoreWithMessages = (messages: any[], sessionId = TEST_IDS.SESSION_ID) => {
+export const setupStoreWithMessages = (
+  messages: any[],
+  sessionId = TEST_IDS.SESSION_ID,
+  topicId: string | null | undefined = TEST_IDS.TOPIC_ID,
+) => {
   useChatStore.setState({
     activeId: sessionId,
+    activeTopicId: topicId ?? undefined,
     messagesMap: {
-      [messageMapKey(sessionId)]: messages,
+      [messageMapKey(sessionId, topicId ?? undefined)]: messages,
     },
   });
 };
@@ -63,6 +68,9 @@ export const spyOnMessageService = () => {
   const updateMessageSpy = vi
     .spyOn(messageService, 'updateMessage')
     .mockResolvedValue({ messages: [], success: true });
+  const updateMessageMetadataSpy = vi
+    .spyOn(messageService, 'updateMessageMetadata')
+    .mockResolvedValue({ messages: [], success: true });
   const removeMessageSpy = vi
     .spyOn(messageService, 'removeMessage')
     .mockResolvedValue(undefined as any);
@@ -74,6 +82,7 @@ export const spyOnMessageService = () => {
     createMessageSpy,
     removeMessageSpy,
     updateMessageErrorSpy,
+    updateMessageMetadataSpy,
     updateMessageSpy,
   };
 };

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

@@ -6,7 +6,7 @@ import { chatService } from '@/services/chat';
 import { ragService } from '@/services/rag';
 import { useAgentStore } from '@/store/agent';
 import { agentSelectors } from '@/store/agent/selectors';
-import { chatSelectors } from '@/store/chat/selectors';
+import { chatSelectors, dbMessageSelectors, displayMessageSelectors } from '@/store/chat/selectors';
 import { systemAgentSelectors } from '@/store/user/selectors';
 import { QueryRewriteSystemAgent } from '@/types/user/settings';
 
@@ -70,7 +70,7 @@ describe('chatRAG actions', () => {
       const userQuery = 'user-query';
 
       // Mock the message with existing ragQuery
-      vi.spyOn(chatSelectors, 'getMessageById').mockReturnValue(
+      vi.spyOn(dbMessageSelectors, 'getDbMessageById').mockReturnValue(
         () =>
           ({
             id: messageId,
@@ -113,7 +113,7 @@ describe('chatRAG actions', () => {
       const rewrittenQuery = 'rewritten-query';
 
       // Mock the message without ragQuery
-      vi.spyOn(chatSelectors, 'getMessageById').mockReturnValue(
+      vi.spyOn(dbMessageSelectors, 'getDbMessageById').mockReturnValue(
         () =>
           ({
             id: messageId,
@@ -220,7 +220,7 @@ describe('chatRAG actions', () => {
     it('should not rewrite if message not found', async () => {
       const { result } = renderHook(() => useChatStore());
 
-      vi.spyOn(chatSelectors, 'getMessageById').mockReturnValue(() => undefined);
+      vi.spyOn(dbMessageSelectors, 'getDbMessageById').mockReturnValue(() => undefined);
       const rewriteSpy = vi.spyOn(result.current, 'internal_rewriteQuery');
 
       await act(async () => {
@@ -235,7 +235,7 @@ describe('chatRAG actions', () => {
       const messageId = 'message-id';
       const content = 'message content';
 
-      vi.spyOn(chatSelectors, 'getMessageById').mockReturnValue(
+      vi.spyOn(dbMessageSelectors, 'getDbMessageById').mockReturnValue(
         () =>
           ({
             id: messageId,
@@ -243,7 +243,7 @@ describe('chatRAG actions', () => {
           }) as UIChatMessage,
       );
 
-      vi.spyOn(chatSelectors, 'mainAIChatsWithHistoryConfig').mockReturnValue([
+      vi.spyOn(displayMessageSelectors, 'mainAIChatsWithHistoryConfig').mockReturnValue([
         { content: 'history' },
       ] as UIChatMessage[]);