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

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

@@ -0,0 +1,258 @@
+import { act, renderHook } from '@testing-library/react';
+import { beforeEach, describe, expect, it } from 'vitest';
+
+import { useChatStore } from '@/store/chat/store';
+
+import { chatToolSelectors } from './selectors';
+
+describe('chatToolSelectors', () => {
+  beforeEach(() => {
+    useChatStore.setState(useChatStore.getInitialState());
+  });
+
+  describe('isDallEImageGenerating', () => {
+    it('should return true when DALL-E image is generating for message', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({
+          dalleImageLoading: {
+            msg1: true,
+            msg2: false,
+          },
+        });
+      });
+
+      expect(chatToolSelectors.isDallEImageGenerating('msg1')(result.current)).toBe(true);
+      expect(chatToolSelectors.isDallEImageGenerating('msg2')(result.current)).toBe(false);
+    });
+
+    it('should return undefined when message not in loading state', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      expect(chatToolSelectors.isDallEImageGenerating('msg1')(result.current)).toBeUndefined();
+    });
+  });
+
+  describe('isGeneratingDallEImage', () => {
+    it('should return true when any DALL-E image is generating', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({
+          dalleImageLoading: {
+            msg1: false,
+            msg2: true,
+            msg3: false,
+          },
+        });
+      });
+
+      expect(chatToolSelectors.isGeneratingDallEImage(result.current)).toBe(true);
+    });
+
+    it('should return false when no DALL-E images are generating', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        useChatStore.setState({
+          dalleImageLoading: {
+            msg1: false,
+            msg2: false,
+          },
+        });
+      });
+
+      expect(chatToolSelectors.isGeneratingDallEImage(result.current)).toBe(false);
+    });
+
+    it('should return false when dalleImageLoading is empty', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      expect(chatToolSelectors.isGeneratingDallEImage(result.current)).toBe(false);
+    });
+  });
+
+  describe('isInterpreterExecuting', () => {
+    it('should return true when interpreter is executing for message', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let opId: string;
+
+      act(() => {
+        opId = result.current.startOperation({
+          type: 'builtinToolInterpreter',
+          context: { sessionId: 'session1', messageId: 'msg1' },
+        }).operationId;
+
+        result.current.associateMessageWithOperation('msg1', opId);
+      });
+
+      expect(chatToolSelectors.isInterpreterExecuting('msg1')(result.current)).toBe(true);
+    });
+
+    it('should return false when no operation exists for message', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      expect(chatToolSelectors.isInterpreterExecuting('msg1')(result.current)).toBe(false);
+    });
+
+    it('should return false when operation is not builtinToolInterpreter', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let opId: string;
+
+      act(() => {
+        opId = result.current.startOperation({
+          type: 'execAgentRuntime',
+          context: { sessionId: 'session1', messageId: 'msg1' },
+        }).operationId;
+
+        result.current.associateMessageWithOperation('msg1', opId);
+      });
+
+      expect(chatToolSelectors.isInterpreterExecuting('msg1')(result.current)).toBe(false);
+    });
+
+    it('should return false when operation is completed', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let opId: string;
+
+      act(() => {
+        opId = result.current.startOperation({
+          type: 'builtinToolInterpreter',
+          context: { sessionId: 'session1', messageId: 'msg1' },
+        }).operationId;
+
+        result.current.associateMessageWithOperation('msg1', opId);
+      });
+
+      act(() => {
+        result.current.completeOperation(opId);
+      });
+
+      expect(chatToolSelectors.isInterpreterExecuting('msg1')(result.current)).toBe(false);
+    });
+  });
+
+  describe('isSearXNGSearching', () => {
+    it('should return true when SearXNG search is running for message', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let opId: string;
+
+      act(() => {
+        opId = result.current.startOperation({
+          type: 'builtinToolSearch',
+          context: { sessionId: 'session1', messageId: 'msg1' },
+        }).operationId;
+
+        result.current.associateMessageWithOperation('msg1', opId);
+      });
+
+      expect(chatToolSelectors.isSearXNGSearching('msg1')(result.current)).toBe(true);
+    });
+
+    it('should return false when no operation exists', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      expect(chatToolSelectors.isSearXNGSearching('msg1')(result.current)).toBe(false);
+    });
+
+    it('should return false when operation type is different', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let opId: string;
+
+      act(() => {
+        opId = result.current.startOperation({
+          type: 'builtinToolInterpreter',
+          context: { sessionId: 'session1', messageId: 'msg1' },
+        }).operationId;
+
+        result.current.associateMessageWithOperation('msg1', opId);
+      });
+
+      expect(chatToolSelectors.isSearXNGSearching('msg1')(result.current)).toBe(false);
+    });
+
+    it('should return false when operation is not running', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let opId: string;
+
+      act(() => {
+        opId = result.current.startOperation({
+          type: 'builtinToolSearch',
+          context: { sessionId: 'session1', messageId: 'msg1' },
+        }).operationId;
+
+        result.current.associateMessageWithOperation('msg1', opId);
+        result.current.completeOperation(opId);
+      });
+
+      expect(chatToolSelectors.isSearXNGSearching('msg1')(result.current)).toBe(false);
+    });
+  });
+
+  describe('isSearchingLocalFiles', () => {
+    it('should return true when local system search is running', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let opId: string;
+
+      act(() => {
+        opId = result.current.startOperation({
+          type: 'builtinToolLocalSystem',
+          context: { sessionId: 'session1', messageId: 'msg1' },
+        }).operationId;
+
+        result.current.associateMessageWithOperation('msg1', opId);
+      });
+
+      expect(chatToolSelectors.isSearchingLocalFiles('msg1')(result.current)).toBe(true);
+    });
+
+    it('should return false when no operation exists', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      expect(chatToolSelectors.isSearchingLocalFiles('msg1')(result.current)).toBe(false);
+    });
+
+    it('should return false when operation type is different', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let opId: string;
+
+      act(() => {
+        opId = result.current.startOperation({
+          type: 'builtinToolSearch',
+          context: { sessionId: 'session1', messageId: 'msg1' },
+        }).operationId;
+
+        result.current.associateMessageWithOperation('msg1', opId);
+      });
+
+      expect(chatToolSelectors.isSearchingLocalFiles('msg1')(result.current)).toBe(false);
+    });
+
+    it('should return false when operation is completed', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      let opId: string;
+
+      act(() => {
+        opId = result.current.startOperation({
+          type: 'builtinToolLocalSystem',
+          context: { sessionId: 'session1', messageId: 'msg1' },
+        }).operationId;
+
+        result.current.associateMessageWithOperation('msg1', opId);
+        result.current.completeOperation(opId);
+      });
+
+      expect(chatToolSelectors.isSearchingLocalFiles('msg1')(result.current)).toBe(false);
+    });
+  });
+});

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

@@ -5,11 +5,32 @@ const isDallEImageGenerating = (id: string) => (s: ChatStoreState) => s.dalleIma
 const isGeneratingDallEImage = (s: ChatStoreState) =>
   Object.values(s.dalleImageLoading).some(Boolean);
 
-const isInterpreterExecuting = (id: string) => (s: ChatStoreState) =>
-  s.codeInterpreterExecuting[id];
+const isInterpreterExecuting = (id: string) => (s: ChatStoreState) => {
+  // Check if there's a running builtinToolInterpreter operation for this message
+  const operationId = s.messageOperationMap[id];
+  if (!operationId) return false;
 
-const isSearXNGSearching = (id: string) => (s: ChatStoreState) => s.searchLoading[id];
-const isSearchingLocalFiles = (id: string) => (s: ChatStoreState) => s.localFileLoading[id];
+  const operation = s.operations[operationId];
+  return operation?.type === 'builtinToolInterpreter' && operation?.status === 'running';
+};
+
+const isSearXNGSearching = (id: string) => (s: ChatStoreState) => {
+  // Check if there's a running builtinToolSearch operation for this message
+  const operationId = s.messageOperationMap[id];
+  if (!operationId) return false;
+
+  const operation = s.operations[operationId];
+  return operation?.type === 'builtinToolSearch' && operation?.status === 'running';
+};
+
+const isSearchingLocalFiles = (id: string) => (s: ChatStoreState) => {
+  // Check if there's a running builtinToolLocalSystem operation for this message
+  const operationId = s.messageOperationMap[id];
+  if (!operationId) return false;
+
+  const operation = s.operations[operationId];
+  return operation?.type === 'builtinToolLocalSystem' && operation?.status === 'running';
+};
 
 export const chatToolSelectors = {
   isDallEImageGenerating,

package/src/store/chat/slices/message/action.test.ts

@@ -28,6 +28,7 @@ vi.mock('@/services/message', () => ({
     createMessage: vi.fn(() => Promise.resolve({ id: 'new-message-id', messages: [] })),
     updateMessage: vi.fn(() => Promise.resolve({ success: true, messages: [] })),
     updateMessageMetadata: vi.fn(() => Promise.resolve({ success: true, messages: [] })),
+    updateMessagePlugin: vi.fn(() => Promise.resolve({ success: true, messages: [] })),
     updateMessagePluginError: vi.fn(() => Promise.resolve({ success: true, messages: [] })),
     updateMessageRAG: vi.fn(() => Promise.resolve({ success: true, messages: [] })),
     removeAllMessages: vi.fn(() => Promise.resolve()),
@@ -688,11 +689,14 @@ describe('chatMessage actions', () => {
         await result.current.optimisticUpdateMessageContent(messageId, newContent);
       });
 
-      expect(internal_dispatchMessageSpy).toHaveBeenCalledWith({
-        id: messageId,
-        type: 'updateMessage',
-        value: { content: newContent },
-      });
+      expect(internal_dispatchMessageSpy).toHaveBeenCalledWith(
+        {
+          id: messageId,
+          type: 'updateMessage',
+          value: { content: newContent },
+        },
+        undefined,
+      );
     });
 
     it('should replace messages after updating content', async () => {
@@ -872,10 +876,17 @@ describe('chatMessage actions', () => {
 
       const updateMessageSpy = vi.spyOn(messageService, 'updateMessage');
 
+      let operationId: string;
       await act(async () => {
+        // Create operation with desired context
+        const op = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: contextSessionId, topicId: contextTopicId },
+        });
+        operationId = op.operationId;
+
         await result.current.optimisticUpdateMessageContent(messageId, content, undefined, {
-          sessionId: contextSessionId,
-          topicId: contextTopicId,
+          operationId,
         });
       });
 
@@ -886,7 +897,7 @@ describe('chatMessage actions', () => {
       );
     });
 
-    it('optimisticUpdateMessageError should use context sessionId/topicId', async () => {
+    it('optimisticUpdateMessageError should use context operationId', async () => {
       const { result } = renderHook(() => useChatStore());
       const messageId = 'message-id';
       const error = { message: 'Error occurred', type: 'error' as any };
@@ -895,10 +906,17 @@ describe('chatMessage actions', () => {
 
       const updateMessageSpy = vi.spyOn(messageService, 'updateMessage');
 
+      let operationId: string;
       await act(async () => {
+        // Create operation with desired context
+        const op = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: contextSessionId, topicId: contextTopicId },
+        });
+        operationId = op.operationId;
+
         await result.current.optimisticUpdateMessageError(messageId, error, {
-          sessionId: contextSessionId,
-          topicId: contextTopicId,
+          operationId,
         });
       });
 
@@ -909,7 +927,7 @@ describe('chatMessage actions', () => {
       );
     });
 
-    it('optimisticDeleteMessage should use context sessionId/topicId', async () => {
+    it('optimisticDeleteMessage should use context operationId', async () => {
       const { result } = renderHook(() => useChatStore());
       const messageId = 'message-id';
       const contextSessionId = 'context-session';
@@ -917,10 +935,17 @@ describe('chatMessage actions', () => {
 
       const removeMessageSpy = vi.spyOn(messageService, 'removeMessage');
 
+      let operationId: string;
       await act(async () => {
+        // Create operation with desired context
+        const op = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: contextSessionId, topicId: contextTopicId },
+        });
+        operationId = op.operationId;
+
         await result.current.optimisticDeleteMessage(messageId, {
-          sessionId: contextSessionId,
-          topicId: contextTopicId,
+          operationId,
         });
       });
 
@@ -930,7 +955,7 @@ describe('chatMessage actions', () => {
       });
     });
 
-    it('optimisticDeleteMessages should use context sessionId/topicId', async () => {
+    it('optimisticDeleteMessages should use context operationId', async () => {
       const { result } = renderHook(() => useChatStore());
       const ids = ['id-1', 'id-2'];
       const contextSessionId = 'context-session';
@@ -938,10 +963,17 @@ describe('chatMessage actions', () => {
 
       const removeMessagesSpy = vi.spyOn(messageService, 'removeMessages');
 
+      let operationId: string;
       await act(async () => {
+        // Create operation with desired context
+        const op = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: contextSessionId, topicId: contextTopicId },
+        });
+        operationId = op.operationId;
+
         await result.current.optimisticDeleteMessages(ids, {
-          sessionId: contextSessionId,
-          topicId: contextTopicId,
+          operationId,
         });
       });
 
@@ -951,4 +983,90 @@ describe('chatMessage actions', () => {
       });
     });
   });
+
+  describe('optimisticUpdateMessagePlugin', () => {
+    it('should dispatch message update action with plugin value', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+      const pluginValue = { arguments: '{"test":"value"}' };
+      const internal_dispatchMessageSpy = vi.spyOn(result.current, 'internal_dispatchMessage');
+
+      await act(async () => {
+        await result.current.optimisticUpdateMessagePlugin(messageId, pluginValue);
+      });
+
+      expect(internal_dispatchMessageSpy).toHaveBeenCalledWith(
+        {
+          id: messageId,
+          type: 'updateMessagePlugin',
+          value: pluginValue,
+        },
+        undefined,
+      );
+    });
+
+    it('should call messageService.updateMessagePlugin with correct parameters', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+      const pluginValue = { state: 'success' };
+
+      const updateMessagePluginSpy = vi.spyOn(messageService, 'updateMessagePlugin');
+      await act(async () => {
+        await result.current.optimisticUpdateMessagePlugin(messageId, pluginValue);
+      });
+
+      expect(updateMessagePluginSpy).toHaveBeenCalledWith(messageId, pluginValue, {
+        sessionId: 'session-id',
+        topicId: 'topic-id',
+      });
+    });
+
+    it('should replace messages after updating plugin', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+      const pluginValue = { apiName: 'test-api' };
+      const replaceMessagesSpy = vi.spyOn(result.current, 'replaceMessages');
+
+      await act(async () => {
+        await result.current.optimisticUpdateMessagePlugin(messageId, pluginValue);
+      });
+
+      expect(replaceMessagesSpy).toHaveBeenCalledWith(
+        [],
+        expect.objectContaining({
+          sessionId: 'session-id',
+          topicId: 'topic-id',
+        }),
+      );
+    });
+
+    it('should use context operationId when provided', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+      const pluginValue = { identifier: 'test-plugin' };
+      const contextSessionId = 'context-session';
+      const contextTopicId = 'context-topic';
+
+      const updateMessagePluginSpy = vi.spyOn(messageService, 'updateMessagePlugin');
+
+      let operationId: string;
+      await act(async () => {
+        // Create operation with desired context
+        const op = result.current.startOperation({
+          type: 'sendMessage',
+          context: { sessionId: contextSessionId, topicId: contextTopicId },
+        });
+        operationId = op.operationId;
+
+        await result.current.optimisticUpdateMessagePlugin(messageId, pluginValue, {
+          operationId,
+        });
+      });
+
+      expect(updateMessagePluginSpy).toHaveBeenCalledWith(messageId, pluginValue, {
+        sessionId: contextSessionId,
+        topicId: contextTopicId,
+      });
+    });
+  });
 });

package/src/store/chat/slices/message/actions/internals.ts

@@ -1,5 +1,6 @@
 import { parse } from '@lobechat/conversation-flow';
 import { TraceEventPayloads } from '@lobechat/types';
+import debug from 'debug';
 import isEqual from 'fast-deep-equal';
 import { StateCreator } from 'zustand/vanilla';
 
@@ -10,6 +11,8 @@ import { displayMessageSelectors } from '../../../selectors';
 import { messageMapKey } from '../../../utils/messageMapKey';
 import { MessageDispatch, messagesReducer } from '../reducer';
 
+const log = debug('lobe-store:message-internals');
+
 /**
  * Internal core methods that serve as building blocks for other actions
  */
@@ -18,10 +21,7 @@ export interface MessageInternalsAction {
    * update message at the frontend
    * this method will not update messages to database
    */
-  internal_dispatchMessage: (
-    payload: MessageDispatch,
-    context?: { sessionId: string; topicId?: string | null },
-  ) => void;
+  internal_dispatchMessage: (payload: MessageDispatch, context?: { operationId?: string }) => void;
 
   /**
    * trace message events for analytics
@@ -37,10 +37,36 @@ export const messageInternals: StateCreator<
 > = (set, get) => ({
   // the internal process method of the AI message
   internal_dispatchMessage: (payload, context) => {
-    const activeId = typeof context !== 'undefined' ? context.sessionId : get().activeId;
-    const topicId = typeof context !== 'undefined' ? context.topicId : get().activeTopicId;
+    let sessionId: string;
+    let topicId: string | null | undefined;
+
+    // Get context from operation if operationId is provided
+    if (context?.operationId) {
+      const operation = get().operations[context.operationId];
+      if (!operation) {
+        log('[internal_dispatchMessage] ERROR: Operation not found: %s', context.operationId);
+        throw new Error(`Operation not found: ${context.operationId}`);
+      }
+      sessionId = operation.context.sessionId!;
+      topicId = operation.context.topicId;
+      log(
+        '[internal_dispatchMessage] get context from operation %s: sessionId=%s, topicId=%s',
+        context.operationId,
+        sessionId,
+        topicId,
+      );
+    } else {
+      // Fallback to global state
+      sessionId = get().activeId;
+      topicId = get().activeTopicId;
+      log(
+        '[internal_dispatchMessage] use global context: sessionId=%s, topicId=%s',
+        sessionId,
+        topicId,
+      );
+    }
 
-    const messagesKey = messageMapKey(activeId, topicId);
+    const messagesKey = messageMapKey(sessionId, topicId);
 
     // Get raw messages from dbMessagesMap and apply reducer
     const rawMessages = get().dbMessagesMap[messagesKey] || [];