@lobehub/lobehub 2.0.0-next.85 → 2.0.0-next.86

package/src/store/chat/slices/aiChat/actions/__tests__/cancel-functionality.test.ts

@@ -1,7 +1,8 @@
 import { act, renderHook } from '@testing-library/react';
-import { describe, expect, it, vi } from 'vitest';
+import { describe, expect, it } from 'vitest';
 
 import { useChatStore } from '../../../../store';
+import { messageMapKey } from '../../../../utils/messageMapKey';
 
 describe('Cancel send message functionality tests', () => {
   describe('cancelSendMessageInServer', () => {
@@ -13,7 +14,8 @@ describe('Cancel send message functionality tests', () => {
         useChatStore.setState({
           activeId: 'session-1',
           activeTopicId: 'topic-1',
-          mainSendMessageOperations: {},
+          operations: {},
+          operationsByContext: {},
         });
       });
 
@@ -28,13 +30,91 @@ describe('Cancel send message functionality tests', () => {
       }).not.toThrow();
     });
 
+    it('should cancel running sendMessage operations', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      const sessionId = 'session-1';
+      const topicId = 'topic-1';
+
+      act(() => {
+        useChatStore.setState({
+          activeId: sessionId,
+          activeTopicId: topicId,
+        });
+      });
+
+      // Start a sendMessage operation
+      let operationId: string;
+      act(() => {
+        const res = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId, topicId },
+        });
+        operationId = res.operationId;
+      });
+
+      expect(result.current.operations[operationId!].status).toBe('running');
+
+      // Cancel the operation
+      act(() => {
+        result.current.cancelSendMessageInServer();
+      });
+
+      expect(result.current.operations[operationId!].status).toBe('cancelled');
+    });
+
+    it('should restore editor state when cancelling', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      const sessionId = 'session-1';
+      const topicId = 'topic-1';
+      const mockEditorState = { content: 'test message' };
+
+      // Mock editor
+      const mockEditor = {
+        setJSONState: vi.fn(),
+        getJSONState: vi.fn().mockReturnValue(mockEditorState),
+      };
+
+      act(() => {
+        useChatStore.setState({
+          activeId: sessionId,
+          activeTopicId: topicId,
+          mainInputEditor: mockEditor as any,
+        });
+      });
+
+      // Create operation with editor state
+      let operationId: string;
+      act(() => {
+        const res = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId, topicId },
+        });
+        operationId = res.operationId;
+
+        result.current.updateOperationMetadata(res.operationId, {
+          inputEditorTempState: mockEditorState,
+        });
+      });
+
+      // Cancel
+      act(() => {
+        result.current.cancelSendMessageInServer();
+      });
+
+      // Verify editor state was restored
+      expect(mockEditor.setJSONState).toHaveBeenCalledWith(mockEditorState);
+    });
+
     it('should be able to call with specified topic ID', () => {
       const { result } = renderHook(() => useChatStore());
 
       act(() => {
         useChatStore.setState({
           activeId: 'session-1',
-          mainSendMessageOperations: {},
+          operations: {},
+          operationsByContext: {},
         });
       });
 
@@ -54,7 +134,8 @@ describe('Cancel send message functionality tests', () => {
         useChatStore.setState({
           activeId: 'session-1',
           activeTopicId: 'topic-1',
-          mainSendMessageOperations: {},
+          operations: {},
+          operationsByContext: {},
         });
       });
 
@@ -66,42 +147,71 @@ describe('Cancel send message functionality tests', () => {
         });
       }).not.toThrow();
     });
-  });
 
-  describe('Internal methods', () => {
-    it('should have internal state management methods', () => {
+    it('should clear error messages from sendMessage operations', () => {
       const { result } = renderHook(() => useChatStore());
 
-      expect(typeof result.current.internal_toggleSendMessageOperation).toBe('function');
-      expect(typeof result.current.internal_updateSendMessageOperation).toBe('function');
-    });
-
-    it('internal_toggleSendMessageOperation should work normally', () => {
-      const { result } = renderHook(() => useChatStore());
+      const sessionId = 'session-1';
+      const topicId = 'topic-1';
 
       act(() => {
-        useChatStore.setState({ mainSendMessageOperations: {} });
+        useChatStore.setState({
+          activeId: sessionId,
+          activeTopicId: topicId,
+        });
       });
 
-      expect(() => {
-        act(() => {
-          const abortController = result.current.internal_toggleSendMessageOperation(
-            'test-key',
-            true,
-          );
-          expect(abortController).toBeInstanceOf(AbortController);
+      // Create operation with error
+      let operationId: string;
+      act(() => {
+        const res = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId, topicId },
         });
-      }).not.toThrow();
+        operationId = res.operationId;
+
+        result.current.updateOperationMetadata(res.operationId, {
+          inputSendErrorMsg: 'Test error',
+        });
+      });
+
+      expect(result.current.operations[operationId!].metadata.inputSendErrorMsg).toBe('Test error');
+
+      // Clear error
+      act(() => {
+        result.current.clearSendMessageError();
+      });
+
+      expect(result.current.operations[operationId!].metadata.inputSendErrorMsg).toBeUndefined();
     });
   });
 
-  describe('State structure', () => {
-    it('should have mainSendMessageOperations state', () => {
+  describe('Operation system', () => {
+    it('should have operation management methods', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      expect(typeof result.current.startOperation).toBe('function');
+      expect(typeof result.current.cancelOperation).toBe('function');
+      expect(typeof result.current.updateOperationMetadata).toBe('function');
+    });
+
+    it('should track operations by context', () => {
       const { result } = renderHook(() => useChatStore());
 
-      // Ensure state exists
-      expect(result.current.mainSendMessageOperations).toBeDefined();
-      expect(typeof result.current.mainSendMessageOperations).toBe('object');
+      const sessionId = 'session-1';
+      const topicId = 'topic-1';
+
+      let operationId: string;
+      act(() => {
+        const res = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId, topicId },
+        });
+        operationId = res.operationId;
+      });
+
+      const contextKey = messageMapKey(sessionId, topicId);
+      expect(result.current.operationsByContext[contextKey]).toContain(operationId!);
     });
   });
 });

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

@@ -19,103 +19,156 @@ afterEach(() => {
 
 describe('ConversationControl actions', () => {
   describe('stopGenerateMessage', () => {
-    it('should abort generation and clear loading state when controller exists', () => {
-      const abortController = new AbortController();
+    it('should cancel running generateAI operations in current context', () => {
+      const { result } = renderHook(() => useChatStore());
 
       act(() => {
-        useChatStore.setState({ chatLoadingIdsAbortController: abortController });
+        useChatStore.setState({
+          activeId: TEST_IDS.SESSION_ID,
+          activeTopicId: TEST_IDS.TOPIC_ID,
+        });
       });
 
-      const { result } = renderHook(() => useChatStore());
-      const toggleLoadingSpy = vi.spyOn(result.current, 'internal_toggleChatLoading');
+      // Create a generateAI operation
+      let operationId: string;
+      act(() => {
+        const res = result.current.startOperation({
+          type: 'execAgentRuntime',
+          context: {
+            sessionId: TEST_IDS.SESSION_ID,
+            topicId: TEST_IDS.TOPIC_ID,
+          },
+        });
+        operationId = res.operationId;
+      });
+
+      expect(result.current.operations[operationId!].status).toBe('running');
 
+      // Stop generation
       act(() => {
         result.current.stopGenerateMessage();
       });
 
-      expect(abortController.signal.aborted).toBe(true);
-      expect(toggleLoadingSpy).toHaveBeenCalledWith(false, undefined, expect.any(String));
+      expect(result.current.operations[operationId!].status).toBe('cancelled');
     });
 
-    it('should do nothing when abort controller is not set', () => {
+    it('should not cancel operations from different context', () => {
+      const { result } = renderHook(() => useChatStore());
+
       act(() => {
-        useChatStore.setState({ chatLoadingIdsAbortController: undefined });
+        useChatStore.setState({
+          activeId: TEST_IDS.SESSION_ID,
+          activeTopicId: TEST_IDS.TOPIC_ID,
+        });
       });
 
-      const { result } = renderHook(() => useChatStore());
-      const toggleLoadingSpy = vi.spyOn(result.current, 'internal_toggleChatLoading');
+      // Create a generateAI operation in a different context
+      let operationId: string;
+      act(() => {
+        const res = result.current.startOperation({
+          type: 'execAgentRuntime',
+          context: {
+            sessionId: 'different-session',
+            topicId: 'different-topic',
+          },
+        });
+        operationId = res.operationId;
+      });
+
+      expect(result.current.operations[operationId!].status).toBe('running');
 
+      // Stop generation - should not affect different context
       act(() => {
         result.current.stopGenerateMessage();
       });
 
-      expect(toggleLoadingSpy).not.toHaveBeenCalled();
+      expect(result.current.operations[operationId!].status).toBe('running');
     });
   });
 
   describe('cancelSendMessageInServer', () => {
-    it('should abort operation and restore editor state when cancelling', () => {
+    it('should cancel operation and restore editor state', () => {
       const { result } = renderHook(() => useChatStore());
-      const mockAbort = vi.fn();
       const mockSetJSONState = vi.fn();
+      const editorState = { content: 'saved content' };
 
       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,
         });
       });
 
+      // Create operation
+      let operationId: string;
+      act(() => {
+        const res = result.current.startOperation({
+          type: 'sendMessage',
+          context: {
+            sessionId: TEST_IDS.SESSION_ID,
+            topicId: TEST_IDS.TOPIC_ID,
+          },
+        });
+        operationId = res.operationId;
+
+        result.current.updateOperationMetadata(res.operationId, {
+          inputEditorTempState: editorState,
+        });
+      });
+
+      expect(result.current.operations[operationId!].status).toBe('running');
+
+      // Cancel
       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' });
+      expect(result.current.operations[operationId!].status).toBe('cancelled');
+      expect(mockSetJSONState).toHaveBeenCalledWith(editorState);
     });
 
     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 },
-            },
+        });
+      });
+
+      // Create operation
+      let operationId: string;
+      act(() => {
+        const res = result.current.startOperation({
+          type: 'sendMessage',
+          context: {
+            sessionId: TEST_IDS.SESSION_ID,
+            topicId: customTopicId,
           },
         });
+        operationId = res.operationId;
       });
 
+      expect(result.current.operations[operationId!].status).toBe('running');
+
+      // Cancel
       act(() => {
         result.current.cancelSendMessageInServer(customTopicId);
       });
 
-      expect(mockAbort).toHaveBeenCalledWith('User cancelled sendMessage operation');
+      expect(result.current.operations[operationId!].status).toBe('cancelled');
     });
 
     it('should handle gracefully when operation does not exist', () => {
       const { result } = renderHook(() => useChatStore());
 
       act(() => {
-        useChatStore.setState({ mainSendMessageOperations: {} });
+        useChatStore.setState({
+          operations: {},
+          operationsByContext: {},
+        });
       });
 
       expect(() => {
@@ -134,31 +187,44 @@ describe('ConversationControl actions', () => {
         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',
-            },
+        });
+      });
+
+      // Create operation with error
+      let operationId: string;
+      act(() => {
+        const res = result.current.startOperation({
+          type: 'sendMessage',
+          context: {
+            sessionId: TEST_IDS.SESSION_ID,
+            topicId: TEST_IDS.TOPIC_ID,
           },
         });
+        operationId = res.operationId;
+
+        result.current.updateOperationMetadata(res.operationId, {
+          inputSendErrorMsg: 'Some error',
+        });
       });
 
+      expect(result.current.operations[operationId!].metadata.inputSendErrorMsg).toBe('Some error');
+
+      // Clear error
       act(() => {
         result.current.clearSendMessageError();
       });
 
-      expect(
-        result.current.mainSendMessageOperations[
-          messageMapKey(TEST_IDS.SESSION_ID, TEST_IDS.TOPIC_ID)
-        ],
-      ).toBeUndefined();
+      expect(result.current.operations[operationId!].metadata.inputSendErrorMsg).toBeUndefined();
     });
 
     it('should handle gracefully when no error operation exists', () => {
       const { result } = renderHook(() => useChatStore());
 
       act(() => {
-        useChatStore.setState({ mainSendMessageOperations: {} });
+        useChatStore.setState({
+          operations: {},
+          operationsByContext: {},
+        });
       });
 
       expect(() => {
@@ -169,112 +235,80 @@ describe('ConversationControl actions', () => {
     });
   });
 
-  describe('internal_toggleSendMessageOperation', () => {
-    it('should create new send operation with abort controller', () => {
+  describe('Operation system integration', () => {
+    it('should create operation with abort controller', () => {
       const { result } = renderHook(() => useChatStore());
+
+      let operationId: string = '';
       let abortController: AbortController | undefined;
 
       act(() => {
-        abortController = result.current.internal_toggleSendMessageOperation('test-key', true);
+        const res = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'test-session' },
+        });
+        operationId = res.operationId;
+        abortController = res.abortController;
       });
 
       expect(abortController!).toBeInstanceOf(AbortController);
-      expect(result.current.mainSendMessageOperations['test-key']?.isLoading).toBe(true);
-      expect(result.current.mainSendMessageOperations['test-key']?.abortController).toBe(
-        abortController,
-      );
+      expect(result.current.operations[operationId!].abortController).toBe(abortController);
+      expect(result.current.operations[operationId!].status).toBe('running');
     });
 
-    it('should stop send operation and clear abort controller', () => {
+    it('should update operation metadata', () => {
       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;
+      let operationId: string;
 
       act(() => {
-        result.current.internal_updateSendMessageOperation('test-key', {
-          isLoading: true,
-          abortController: mockAbortController,
+        const res = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'test-session' },
         });
+        operationId = res.operationId;
 
-        result.current.internal_toggleSendMessageOperation('test-key', false, 'Test cancel reason');
+        result.current.updateOperationMetadata(res.operationId, {
+          inputSendErrorMsg: 'test error',
+          inputEditorTempState: { content: 'test' },
+        });
       });
 
-      expect(mockAbortController.abort).toHaveBeenCalledWith('Test cancel reason');
+      expect(result.current.operations[operationId!].metadata.inputSendErrorMsg).toBe('test error');
+      expect(result.current.operations[operationId!].metadata.inputEditorTempState).toEqual({
+        content: 'test',
+      });
     });
 
     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();
+      let opId1: string = '';
+      let opId2: string = '';
 
       act(() => {
-        result.current.internal_updateSendMessageOperation('test-key', {
-          isLoading: true,
-          abortController: mockAbortController,
-          inputSendErrorMsg: 'test error',
+        const res1 = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'session-1', topicId: 'topic-1' },
+        });
+        const res2 = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'session-1', topicId: 'topic-2' },
         });
-      });
 
-      expect(result.current.mainSendMessageOperations['test-key']).toEqual({
-        isLoading: true,
-        abortController: mockAbortController,
-        inputSendErrorMsg: 'test error',
+        opId1 = res1.operationId;
+        opId2 = res2.operationId;
       });
-    });
-
-    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,
-        });
+      expect(result.current.operations[opId1!].status).toBe('running');
+      expect(result.current.operations[opId2!].status).toBe('running');
+      expect(opId1).not.toBe(opId2);
 
-        result.current.internal_updateSendMessageOperation('test-key', {
-          inputSendErrorMsg: 'new error',
-        });
-      });
+      const contextKey1 = messageMapKey('session-1', 'topic-1');
+      const contextKey2 = messageMapKey('session-1', 'topic-2');
 
-      expect(result.current.mainSendMessageOperations['test-key']).toEqual({
-        isLoading: true,
-        abortController: initialController,
-        inputSendErrorMsg: 'new error',
-      });
+      expect(result.current.operationsByContext[contextKey1]).toContain(opId1!);
+      expect(result.current.operationsByContext[contextKey2]).toContain(opId2!);
     });
   });
 
@@ -292,7 +326,9 @@ describe('ConversationControl actions', () => {
         await result.current.switchMessageBranch(messageId, branchIndex);
       });
 
-      expect(optimisticUpdateSpy).toHaveBeenCalledWith(messageId, { activeBranchIndex: branchIndex });
+      expect(optimisticUpdateSpy).toHaveBeenCalledWith(messageId, {
+        activeBranchIndex: branchIndex,
+      });
     });
 
     it('should handle switching to branch 0', async () => {
@@ -326,7 +362,9 @@ describe('ConversationControl actions', () => {
         }),
       ).rejects.toThrow('Update failed');
 
-      expect(optimisticUpdateSpy).toHaveBeenCalledWith(messageId, { activeBranchIndex: branchIndex });
+      expect(optimisticUpdateSpy).toHaveBeenCalledWith(messageId, {
+        activeBranchIndex: branchIndex,
+      });
     });
   });
 });

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

@@ -151,9 +151,14 @@ describe('ConversationLifecycle actions', () => {
     it('should not regenerate when already regenerating', async () => {
       const { result } = renderHook(() => useChatStore());
 
+      // Create a regenerate operation to simulate already regenerating
       act(() => {
+        const { operationId } = result.current.startOperation({
+          type: 'regenerate',
+          context: { sessionId: TEST_IDS.SESSION_ID, messageId: TEST_IDS.USER_MESSAGE_ID },
+        });
+
         useChatStore.setState({
-          regeneratingIds: [TEST_IDS.USER_MESSAGE_ID],
           internal_execAgentRuntime: vi.fn(),
         });
       });
@@ -204,9 +209,14 @@ describe('ConversationLifecycle actions', () => {
     it('should not regenerate when already regenerating', async () => {
       const { result } = renderHook(() => useChatStore());
 
+      // Create a regenerate operation to simulate already regenerating
       act(() => {
+        result.current.startOperation({
+          type: 'regenerate',
+          context: { sessionId: TEST_IDS.SESSION_ID, messageId: TEST_IDS.MESSAGE_ID },
+        });
+
         useChatStore.setState({
-          regeneratingIds: [TEST_IDS.MESSAGE_ID],
           internal_execAgentRuntime: vi.fn(),
         });
       });

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

@@ -59,8 +59,6 @@ export const createMockChatConfig = (overrides = {}) => ({
 export const createMockStoreState = (overrides = {}) => ({
   activeId: TEST_IDS.SESSION_ID,
   activeTopicId: TEST_IDS.TOPIC_ID,
-  chatLoadingIds: [],
-  chatLoadingIdsAbortController: undefined,
   messagesMap: {},
   toolCallingStreamIds: {},
   ...overrides,