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

package/src/store/chat/slices/aiChat/selectors.test.ts

@@ -0,0 +1,280 @@
+import { act, renderHook } from '@testing-library/react';
+import { beforeEach, describe, expect, it } from 'vitest';
+
+import { useChatStore } from '@/store/chat/store';
+
+import { aiChatSelectors } from './selectors';
+
+describe('aiChatSelectors', () => {
+  beforeEach(() => {
+    useChatStore.setState(useChatStore.getInitialState());
+  });
+
+  describe('isMessageInReasoning', () => {
+    it('should return true when message has reasoning operation', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({ activeId: 'session1', activeTopicId: 'topic1' });
+        result.current.startOperation({
+          type: 'reasoning',
+          context: { sessionId: 'session1', topicId: 'topic1', messageId: 'msg1' },
+        });
+      });
+
+      expect(aiChatSelectors.isMessageInReasoning('msg1')(result.current)).toBe(true);
+    });
+
+    it('should return false when message has no reasoning operation', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      expect(aiChatSelectors.isMessageInReasoning('msg1')(result.current)).toBe(false);
+    });
+  });
+
+  describe('isMessageInSearchWorkflow', () => {
+    it('should return true when message is in search workflow', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({ searchWorkflowLoadingIds: ['msg1', 'msg2'] });
+      });
+
+      expect(aiChatSelectors.isMessageInSearchWorkflow('msg1')(result.current)).toBe(true);
+      expect(aiChatSelectors.isMessageInSearchWorkflow('msg2')(result.current)).toBe(true);
+    });
+
+    it('should return false when message is not in search workflow', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({ searchWorkflowLoadingIds: ['msg1'] });
+      });
+
+      expect(aiChatSelectors.isMessageInSearchWorkflow('msg2')(result.current)).toBe(false);
+    });
+  });
+
+  describe('isIntentUnderstanding', () => {
+    it('should return true when message is in search workflow', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({ searchWorkflowLoadingIds: ['msg1'] });
+      });
+
+      expect(aiChatSelectors.isIntentUnderstanding('msg1')(result.current)).toBe(true);
+    });
+
+    it('should return false when message is not in search workflow', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      expect(aiChatSelectors.isIntentUnderstanding('msg1')(result.current)).toBe(false);
+    });
+  });
+
+  describe('isCurrentSendMessageLoading', () => {
+    it('should return true when there is a running sendMessage operation in current context', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({ activeId: 'session1', activeTopicId: 'topic1' });
+        result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'session1', topicId: 'topic1' },
+        });
+      });
+
+      expect(aiChatSelectors.isCurrentSendMessageLoading(result.current)).toBe(true);
+    });
+
+    it('should return false when there is no sendMessage operation', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({ activeId: 'session1', activeTopicId: 'topic1' });
+      });
+
+      expect(aiChatSelectors.isCurrentSendMessageLoading(result.current)).toBe(false);
+    });
+
+    it('should return false when sendMessage operation is completed', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let opId: string;
+
+      act(() => {
+        useChatStore.setState({ activeId: 'session1', activeTopicId: 'topic1' });
+        opId = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'session1', topicId: 'topic1' },
+        }).operationId;
+      });
+
+      act(() => {
+        result.current.completeOperation(opId);
+      });
+
+      expect(aiChatSelectors.isCurrentSendMessageLoading(result.current)).toBe(false);
+    });
+
+    it('should return false for different context', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({ activeId: 'session1', activeTopicId: 'topic1' });
+        result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'session2', topicId: 'topic2' },
+        });
+      });
+
+      expect(aiChatSelectors.isCurrentSendMessageLoading(result.current)).toBe(false);
+    });
+  });
+
+  describe('isCurrentSendMessageError', () => {
+    it('should return error message when latest sendMessage operation has error', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let opId: string;
+
+      act(() => {
+        useChatStore.setState({ activeId: 'session1', activeTopicId: 'topic1' });
+        opId = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'session1', topicId: 'topic1' },
+        }).operationId;
+      });
+
+      act(() => {
+        result.current.updateOperationMetadata(opId, {
+          inputSendErrorMsg: 'Network error',
+        });
+      });
+
+      expect(aiChatSelectors.isCurrentSendMessageError(result.current)).toBe('Network error');
+    });
+
+    it('should return undefined when there is no error', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({ activeId: 'session1', activeTopicId: 'topic1' });
+        result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'session1', topicId: 'topic1' },
+        });
+      });
+
+      expect(aiChatSelectors.isCurrentSendMessageError(result.current)).toBeUndefined();
+    });
+
+    it('should return undefined when there are no operations', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({ activeId: 'session1', activeTopicId: 'topic1' });
+      });
+
+      expect(aiChatSelectors.isCurrentSendMessageError(result.current)).toBeUndefined();
+    });
+
+    it('should return the latest error when multiple operations exist', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let op1Id: string;
+      let op2Id: string;
+
+      act(() => {
+        useChatStore.setState({ activeId: 'session1', activeTopicId: 'topic1' });
+
+        op1Id = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'session1', topicId: 'topic1' },
+        }).operationId;
+
+        op2Id = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'session1', topicId: 'topic1' },
+        }).operationId;
+      });
+
+      act(() => {
+        result.current.updateOperationMetadata(op1Id, {
+          inputSendErrorMsg: 'First error',
+        });
+        result.current.updateOperationMetadata(op2Id, {
+          inputSendErrorMsg: 'Second error',
+        });
+      });
+
+      // Should return the latest (second) error
+      expect(aiChatSelectors.isCurrentSendMessageError(result.current)).toBe('Second error');
+    });
+  });
+
+  describe('isSendMessageLoadingForTopic', () => {
+    it('should return true when sendMessage operation is running for the topic', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'session1', topicId: 'topic1' },
+        });
+      });
+
+      expect(aiChatSelectors.isSendMessageLoadingForTopic('session1_topic1')(result.current)).toBe(
+        true,
+      );
+    });
+
+    it('should return false when no sendMessage operation exists', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      expect(aiChatSelectors.isSendMessageLoadingForTopic('session1_topic1')(result.current)).toBe(
+        false,
+      );
+    });
+
+    it('should return false when sendMessage operation is completed', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let opId: string;
+
+      act(() => {
+        opId = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'session1', topicId: 'topic1' },
+        }).operationId;
+      });
+
+      act(() => {
+        result.current.completeOperation(opId);
+      });
+
+      expect(aiChatSelectors.isSendMessageLoadingForTopic('session1_topic1')(result.current)).toBe(
+        false,
+      );
+    });
+
+    it('should distinguish between different topics', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: 'session1', topicId: 'topic1' },
+        });
+      });
+
+      expect(aiChatSelectors.isSendMessageLoadingForTopic('session1_topic1')(result.current)).toBe(
+        true,
+      );
+      expect(aiChatSelectors.isSendMessageLoadingForTopic('session1_topic2')(result.current)).toBe(
+        false,
+      );
+    });
+  });
+});

package/src/store/chat/slices/aiChat/selectors.ts

@@ -1,9 +1,10 @@
 import { messageMapKey } from '@/store/chat/utils/messageMapKey';
 
 import type { ChatStoreState } from '../../initialState';
+import { operationSelectors } from '../operation/selectors';
 
 const isMessageInReasoning = (id: string) => (s: ChatStoreState) =>
-  s.reasoningLoadingIds.includes(id);
+  operationSelectors.isMessageInReasoning(id)(s);
 
 const isMessageInSearchWorkflow = (id: string) => (s: ChatStoreState) =>
   s.searchWorkflowLoadingIds.includes(id);
@@ -12,17 +13,40 @@ const isIntentUnderstanding = (id: string) => (s: ChatStoreState) =>
   isMessageInSearchWorkflow(id)(s);
 
 const isCurrentSendMessageLoading = (s: ChatStoreState) => {
-  const operationKey = messageMapKey(s.activeId, s.activeTopicId);
-  return s.mainSendMessageOperations[operationKey]?.isLoading || false;
+  const contextKey = messageMapKey(s.activeId, s.activeTopicId);
+  const operationIds = s.operationsByContext[contextKey] || [];
+
+  // Check if there's any running sendMessage operation
+  return operationIds.some((opId) => {
+    const op = s.operations[opId];
+    return op && op.type === 'sendMessage' && op.status === 'running';
+  });
 };
 
 const isCurrentSendMessageError = (s: ChatStoreState) => {
-  const operationKey = messageMapKey(s.activeId, s.activeTopicId);
-  return s.mainSendMessageOperations[operationKey]?.inputSendErrorMsg;
+  const contextKey = messageMapKey(s.activeId, s.activeTopicId);
+  const operationIds = s.operationsByContext[contextKey] || [];
+
+  // Find the latest sendMessage operation with error
+  for (const opId of [...operationIds].reverse()) {
+    const op = s.operations[opId];
+    if (op && op.type === 'sendMessage' && op.metadata.inputSendErrorMsg) {
+      return op.metadata.inputSendErrorMsg;
+    }
+  }
+
+  return undefined;
 };
 
-const isSendMessageLoadingForTopic = (topicKey: string) => (s: ChatStoreState) =>
-  s.mainSendMessageOperations[topicKey]?.isLoading ?? false;
+const isSendMessageLoadingForTopic = (topicKey: string) => (s: ChatStoreState) => {
+  const operationIds = s.operationsByContext[topicKey] || [];
+
+  // Check if there's any running sendMessage operation for this topic
+  return operationIds.some((opId) => {
+    const op = s.operations[opId];
+    return op && op.type === 'sendMessage' && op.status === 'running';
+  });
+};
 
 export const aiChatSelectors = {
   isCurrentSendMessageError,

package/src/store/chat/slices/builtinTool/actions/__tests__/localSystem.test.ts

@@ -21,13 +21,19 @@ vi.mock('@/services/electron/localFileService', () => ({
 const mockSet = vi.fn();
 
 const mockStore = {
+  completeOperation: vi.fn(),
+  failOperation: vi.fn(),
   internal_triggerLocalFileToolCalling: vi.fn(),
+  messageOperationMap: {},
   optimisticUpdateMessageContent: vi.fn(),
   optimisticUpdateMessagePluginError: vi.fn(),
   optimisticUpdatePluginArguments: vi.fn(),
   optimisticUpdatePluginState: vi.fn(),
   set: mockSet,
-  toggleLocalFileLoading: vi.fn(),
+  startOperation: vi.fn().mockReturnValue({
+    abortController: new AbortController(),
+    operationId: 'test-op-id',
+  }),
 } as unknown as ChatStore;
 
 const createStore = () => {
@@ -58,10 +64,10 @@ describe('localFileSlice', () => {
 
       await store.internal_triggerLocalFileToolCalling('test-id', mockService);
 
-      expect(mockStore.toggleLocalFileLoading).toBeCalledWith('test-id', true);
-      expect(mockStore.optimisticUpdatePluginState).toBeCalledWith('test-id', mockState);
-      expect(mockStore.optimisticUpdateMessageContent).toBeCalledWith('test-id', mockContent);
-      expect(mockStore.toggleLocalFileLoading).toBeCalledWith('test-id', false);
+      expect(mockStore.startOperation).toBeCalled();
+      expect(mockStore.completeOperation).toBeCalled();
+      expect(mockStore.optimisticUpdatePluginState).toBeCalled();
+      expect(mockStore.optimisticUpdateMessageContent).toBeCalled();
     });
 
     it('should handle error', async () => {
@@ -70,11 +76,15 @@ describe('localFileSlice', () => {
 
       await store.internal_triggerLocalFileToolCalling('test-id', mockService);
 
-      expect(mockStore.optimisticUpdateMessagePluginError).toBeCalledWith('test-id', {
-        body: mockError,
-        message: 'test error',
-        type: 'PluginServerError',
-      });
+      expect(mockStore.optimisticUpdateMessagePluginError).toBeCalledWith(
+        'test-id',
+        {
+          body: mockError,
+          message: 'test error',
+          type: 'PluginServerError',
+        },
+        { operationId: 'test-op-id' },
+      );
     });
   });
 
@@ -187,24 +197,5 @@ describe('localFileSlice', () => {
     });
   });
 
-  describe('toggleLocalFileLoading', () => {
-    it('should toggle loading state', () => {
-      const mockSetFn = vi.fn();
-      const testStore = localSystemSlice(mockSetFn, () => mockStore, {} as any);
-
-      testStore.toggleLocalFileLoading('test-id', true);
-      expect(mockSetFn).toHaveBeenCalledWith(
-        expect.any(Function),
-        false,
-        'toggleLocalFileLoading/start',
-      );
-
-      testStore.toggleLocalFileLoading('test-id', false);
-      expect(mockSetFn).toHaveBeenCalledWith(
-        expect.any(Function),
-        false,
-        'toggleLocalFileLoading/end',
-      );
-    });
-  });
+  // toggleLocalFileLoading is no longer needed as we use operation-based state management
 });

package/src/store/chat/slices/builtinTool/actions/__tests__/search.test.ts

@@ -85,16 +85,18 @@ describe('search actions', () => {
         },
       ];
 
-      expect(searchService.webSearch).toHaveBeenCalledWith({
-        searchEngines: ['google'],
-        query: 'test query',
-      });
-      expect(result.current.searchLoading[messageId]).toBe(false);
+      expect(searchService.webSearch).toHaveBeenCalledWith(
+        {
+          searchEngines: ['google'],
+          query: 'test query',
+        },
+        { signal: expect.any(AbortSignal) },
+      );
       expect(result.current.optimisticUpdateMessageContent).toHaveBeenCalledWith(
         messageId,
         searchResultsPrompt(expectedContent),
         undefined,
-        { sessionId: undefined, topicId: undefined },
+        { operationId: expect.any(String) },
       );
     });
 
@@ -122,17 +124,19 @@ describe('search actions', () => {
         await search(messageId, query);
       });
 
-      expect(searchService.webSearch).toHaveBeenCalledWith({
-        searchEngines: ['custom-engine'],
-        searchTimeRange: 'year',
-        query: 'test query',
-      });
-      expect(result.current.searchLoading[messageId]).toBe(false);
+      expect(searchService.webSearch).toHaveBeenCalledWith(
+        {
+          searchEngines: ['custom-engine'],
+          searchTimeRange: 'year',
+          query: 'test query',
+        },
+        { signal: expect.any(AbortSignal) },
+      );
       expect(result.current.optimisticUpdateMessageContent).toHaveBeenCalledWith(
         messageId,
         searchResultsPrompt([]),
         undefined,
-        { sessionId: undefined, topicId: undefined },
+        { operationId: expect.any(String) },
       );
     });
 
@@ -159,14 +163,7 @@ describe('search actions', () => {
           message: 'Search failed',
           type: 'PluginServerError',
         },
-        { sessionId: undefined, topicId: undefined },
-      );
-      expect(result.current.searchLoading[messageId]).toBe(false);
-      expect(result.current.optimisticUpdateMessageContent).toHaveBeenCalledWith(
-        messageId,
-        'Search failed',
-        undefined,
-        { sessionId: undefined, topicId: undefined },
+        { operationId: expect.any(String) },
       );
     });
   });
@@ -207,7 +204,7 @@ describe('search actions', () => {
         messageId,
         crawlResultsPrompt(expectedContent as any),
         undefined,
-        { sessionId: undefined, topicId: undefined },
+        { operationId: expect.any(String) },
       );
     });
 
@@ -235,7 +232,7 @@ describe('search actions', () => {
         messageId,
         crawlResultsPrompt(mockResponse.results),
         undefined,
-        { sessionId: undefined, topicId: undefined },
+        { operationId: expect.any(String) },
       );
     });
   });
@@ -304,7 +301,7 @@ describe('search actions', () => {
           sessionId: 'session-id',
           topicId: 'topic-id',
         }),
-        { sessionId: 'session-id', topicId: 'topic-id' },
+        { operationId: expect.any(String) },
       );
 
       expect(result.current.optimisticAddToolToAssistantMessage).toHaveBeenCalledWith(
@@ -313,7 +310,7 @@ describe('search actions', () => {
           identifier: 'search',
           type: 'default',
         }),
-        { sessionId: undefined, topicId: undefined },
+        { operationId: expect.any(String) },
       );
     });
 
@@ -332,24 +329,7 @@ describe('search actions', () => {
     });
   });
 
-  describe('toggleSearchLoading', () => {
-    it('should toggle search loading state', () => {
-      const { result } = renderHook(() => useChatStore());
-      const messageId = 'test-message-id';
-
-      act(() => {
-        result.current.toggleSearchLoading(messageId, true);
-      });
-
-      expect(result.current.searchLoading[messageId]).toBe(true);
-
-      act(() => {
-        result.current.toggleSearchLoading(messageId, false);
-      });
-
-      expect(result.current.searchLoading[messageId]).toBe(false);
-    });
-  });
+  // toggleSearchLoading is no longer needed as we use operation-based state management
 
   describe('OptimisticUpdateContext isolation', () => {
     it('search should pass context to optimistic methods', async () => {
@@ -401,13 +381,13 @@ describe('search actions', () => {
       expect(result.current.optimisticUpdatePluginState).toHaveBeenCalledWith(
         messageId,
         expect.any(Object),
-        { sessionId: contextSessionId, topicId: contextTopicId },
+        { operationId: expect.any(String) },
       );
       expect(result.current.optimisticUpdateMessageContent).toHaveBeenCalledWith(
         messageId,
         expect.any(String),
         undefined,
-        { sessionId: contextSessionId, topicId: contextTopicId },
+        { operationId: expect.any(String) },
       );
     });
 
@@ -456,12 +436,12 @@ describe('search actions', () => {
         messageId,
         expect.any(String),
         undefined,
-        { sessionId: contextSessionId, topicId: contextTopicId },
+        { operationId: expect.any(String) },
       );
       expect(result.current.optimisticUpdatePluginState).toHaveBeenCalledWith(
         messageId,
         expect.any(Object),
-        { sessionId: contextSessionId, topicId: contextTopicId },
+        { operationId: expect.any(String) },
       );
     });
 
@@ -510,14 +490,14 @@ describe('search actions', () => {
           identifier: 'search',
           type: 'default',
         }),
-        { sessionId: contextSessionId, topicId: contextTopicId },
+        { operationId: expect.any(String) },
       );
       expect(result.current.optimisticCreateMessage).toHaveBeenCalledWith(
         expect.objectContaining({
           sessionId: contextSessionId,
           topicId: contextTopicId,
         }),
-        { sessionId: contextSessionId, topicId: contextTopicId },
+        { operationId: expect.any(String) },
       );
     });
   });