@lobehub/chat 1.2.6 → 1.2.7

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

@@ -30,6 +30,7 @@ vi.mock('@/services/message', () => ({
     getMessages: vi.fn(),
     updateMessageError: vi.fn(),
     removeMessage: vi.fn(),
+    removeMessagesByAssistant: vi.fn(),
     removeMessages: vi.fn(() => Promise.resolve()),
     createMessage: vi.fn(() => Promise.resolve('new-message-id')),
     updateMessage: vi.fn(),
@@ -147,6 +148,124 @@ describe('chatMessage actions', () => {
       expect(deleteSpy).toHaveBeenCalledWith(messageId);
       expect(result.current.refreshMessages).toHaveBeenCalled();
     });
+
+    it('deleteMessage should remove messages with tools', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+      const removeMessagesSpy = vi.spyOn(messageService, 'removeMessages');
+
+      act(() => {
+        useChatStore.setState({
+          activeId: 'session-id',
+          activeTopicId: undefined,
+          messagesMap: {
+            [messageMapKey('session-id')]: [
+              { id: messageId, tools: [{ id: 'tool1' }, { id: 'tool2' }] } as ChatMessage,
+              { id: '2', tool_call_id: 'tool1', role: 'tool' } as ChatMessage,
+              { id: '3', tool_call_id: 'tool2', role: 'tool' } as ChatMessage,
+            ],
+          },
+        });
+      });
+      await act(async () => {
+        await result.current.deleteMessage(messageId);
+      });
+
+      expect(removeMessagesSpy).toHaveBeenCalledWith([messageId, '2', '3']);
+      expect(result.current.refreshMessages).toHaveBeenCalled();
+    });
+  });
+
+  describe('deleteToolMessage', () => {
+    it('deleteMessage should remove a message by id', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+      const updateMessageSpy = vi.spyOn(messageService, 'updateMessage');
+      const removeMessageSpy = vi.spyOn(messageService, 'removeMessage');
+
+      act(() => {
+        useChatStore.setState({
+          activeId: 'session-id',
+          activeTopicId: undefined,
+          messagesMap: {
+            [messageMapKey('session-id')]: [
+              {
+                id: messageId,
+                role: 'assistant',
+                tools: [{ id: 'tool1' }, { id: 'tool2' }],
+              } as ChatMessage,
+              { id: '2', parentId: messageId, tool_call_id: 'tool1', role: 'tool' } as ChatMessage,
+              { id: '3', tool_call_id: 'tool2', role: 'tool' } as ChatMessage,
+            ],
+          },
+        });
+      });
+      await act(async () => {
+        await result.current.deleteToolMessage('2');
+      });
+
+      expect(removeMessageSpy).toHaveBeenCalled();
+      expect(updateMessageSpy).toHaveBeenCalledWith('message-id', {
+        tools: [{ id: 'tool2' }],
+      });
+      expect(result.current.refreshMessages).toHaveBeenCalled();
+    });
+  });
+
+  describe('delAndRegenerateMessage', () => {
+    it('should remove a message and create a new message', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+      const deleteMessageSpy = vi.spyOn(result.current, 'deleteMessage');
+      const resendMessageSpy = vi.spyOn(result.current, 'internal_resendMessage');
+
+      act(() => {
+        useChatStore.setState({
+          activeId: 'session-id',
+          activeTopicId: undefined,
+          messagesMap: {
+            [messageMapKey('session-id')]: [
+              { id: messageId, tools: [{ id: 'tool1' }, { id: 'tool2' }] } as ChatMessage,
+            ],
+          },
+        });
+      });
+      await act(async () => {
+        await result.current.delAndRegenerateMessage(messageId);
+      });
+
+      expect(deleteMessageSpy).toHaveBeenCalledWith(messageId);
+      expect(resendMessageSpy).toHaveBeenCalled();
+      expect(result.current.refreshMessages).toHaveBeenCalled();
+    });
+  });
+  describe('regenerateMessage', () => {
+    it('should create a new message', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+      const resendMessageSpy = vi.spyOn(result.current, 'internal_resendMessage');
+
+      act(() => {
+        useChatStore.setState({
+          activeId: 'session-id',
+          activeTopicId: undefined,
+          messagesMap: {
+            [messageMapKey('session-id')]: [
+              {
+                id: messageId,
+                tools: [{ id: 'tool1' }, { id: 'tool2' }],
+                traceId: 'abc',
+              } as ChatMessage,
+            ],
+          },
+        });
+      });
+      await act(async () => {
+        await result.current.regenerateMessage(messageId);
+      });
+
+      expect(resendMessageSpy).toHaveBeenCalledWith(messageId, 'abc');
+    });
   });
 
   describe('clearAllMessages', () => {
@@ -1144,6 +1263,31 @@ describe('chatMessage actions', () => {
     });
   });
 
+  describe('internal_toggleToolCallingStreaming action', () => {
+    it('should add message id to messageLoadingIds when loading is true', () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'message-id';
+
+      act(() => {
+        result.current.internal_toggleToolCallingStreaming(messageId, [true]);
+      });
+
+      expect(result.current.toolCallingStreamIds[messageId]).toEqual([true]);
+    });
+
+    it('should remove message id from messageLoadingIds when loading is false', () => {
+      const { result } = renderHook(() => useChatStore());
+      const messageId = 'ddd-id';
+
+      act(() => {
+        result.current.internal_toggleToolCallingStreaming(messageId, [true]);
+        result.current.internal_toggleToolCallingStreaming(messageId, undefined);
+      });
+
+      expect(result.current.toolCallingStreamIds[messageId]).toBeUndefined();
+    });
+  });
+
   describe('stopGenerateMessage', () => {
     it('should return early if abortController is undefined', () => {
       act(() => {

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

@@ -66,6 +66,7 @@ export interface ChatMessageAction {
    */
   clearMessage: () => Promise<void>;
   deleteMessage: (id: string) => Promise<void>;
+  deleteToolMessage: (id: string) => Promise<void>;
   delAndRegenerateMessage: (id: string) => Promise<void>;
   clearAllMessages: () => Promise<void>;
   // update
@@ -81,19 +82,7 @@ export interface ChatMessageAction {
   // =========  ↓ Internal Method ↓  ========== //
   // ========================================== //
   // ========================================== //
-  internal_toggleChatLoading: (
-    loading: boolean,
-    id?: string,
-    action?: string,
-  ) => AbortController | undefined;
-  internal_toggleLoadingArrays: (
-    key: keyof ChatStoreState,
-    loading: boolean,
-    id?: string,
-    action?: string,
-  ) => AbortController | undefined;
-  internal_toggleToolCallingStreaming: (id: string, streaming: boolean[] | undefined) => void;
-  internal_toggleMessageLoading: (loading: boolean, id: string) => void;
+
   /**
    * update message at the frontend point
    * this method will not update messages to database
@@ -108,9 +97,7 @@ export interface ChatMessageAction {
     params?: ProcessMessageParams,
   ) => Promise<void>;
   /**
-   * 实际获取 AI 响应
-   * @param messages - 聊天消息数组
-   * @param options - 获取 SSE 选项
+   * the method to fetch the AI message
    */
   internal_fetchAIChatMessage: (
     messages: ChatMessage[],
@@ -122,24 +109,66 @@ export interface ChatMessageAction {
   }>;
 
   /**
+   * update the message content with optimistic update
    * a method used by other action
-   * @param id
-   * @param content
    */
   internal_updateMessageContent: (
     id: string,
     content: string,
     toolCalls?: MessageToolCall[],
   ) => Promise<void>;
+  /**
+   * update the message error with optimistic update
+   */
   internal_updateMessageError: (id: string, error: ChatMessageError | null) => Promise<void>;
+  /**
+   * create a message with optimistic update
+   */
   internal_createMessage: (
     params: CreateMessageParams,
     context?: { tempMessageId?: string; skipRefresh?: boolean },
   ) => Promise<string>;
+  /**
+   * create a temp message for optimistic update
+   * otherwise the message will be too slow to show
+   */
   internal_createTmpMessage: (params: CreateMessageParams) => string;
-  internal_fetchMessages: () => Promise<void>;
+  /**
+   * delete the message content with optimistic update
+   */
+  internal_deleteMessage: (id: string) => Promise<void>;
   internal_resendMessage: (id: string, traceId?: string) => Promise<void>;
+
+  internal_fetchMessages: () => Promise<void>;
   internal_traceMessage: (id: string, payload: TraceEventPayloads) => Promise<void>;
+
+  /**
+   * method to toggle message create loading state
+   * the AI message status is creating -> generating
+   * other message role like user and tool , only this method need to be called
+   */
+  internal_toggleMessageLoading: (loading: boolean, id: string) => void;
+  /**
+   * method to toggle ai message generating loading
+   */
+  internal_toggleChatLoading: (
+    loading: boolean,
+    id?: string,
+    action?: string,
+  ) => AbortController | undefined;
+  /**
+   * method to toggle the tool calling loading state
+   */
+  internal_toggleToolCallingStreaming: (id: string, streaming: boolean[] | undefined) => void;
+  /**
+   * helper to toggle the loading state of the array,used by these three toggleXXXLoading
+   */
+  internal_toggleLoadingArrays: (
+    key: keyof ChatStoreState,
+    loading: boolean,
+    id?: string,
+    action?: string,
+  ) => AbortController | undefined;
 }
 
 const getAgentConfig = () => agentSelectors.currentAgentConfig(useAgentStore.getState());
@@ -155,30 +184,42 @@ export const chatMessage: StateCreator<
     const message = chatSelectors.getMessageById(id)(get());
     if (!message) return;
 
-    const deleteFn = async (id: string) => {
-      get().internal_dispatchMessage({ type: 'deleteMessage', id });
-      await messageService.removeMessage(id);
-    };
+    let ids = [message.id];
 
     // if the message is a tool calls, then delete all the related messages
-    // TODO: maybe we need to delete it in the DB?
     if (message.tools) {
-      const pools = message.tools
-        .flatMap((tool) => {
-          const messages = chatSelectors
-            .currentChats(get())
-            .filter((m) => m.tool_call_id === tool.id);
-
-          return messages.map((m) => m.id);
-        })
-        .map((i) => deleteFn(i));
+      const toolMessageIds = message.tools.flatMap((tool) => {
+        const messages = chatSelectors
+          .currentChats(get())
+          .filter((m) => m.tool_call_id === tool.id);
 
-      await Promise.all(pools);
+        return messages.map((m) => m.id);
+      });
+      ids = ids.concat(toolMessageIds);
     }
 
-    await deleteFn(id);
+    get().internal_dispatchMessage({ type: 'deleteMessages', ids });
+    await messageService.removeMessages(ids);
     await get().refreshMessages();
   },
+
+  deleteToolMessage: async (id) => {
+    const message = chatSelectors.getMessageById(id)(get());
+    if (!message || message.role !== 'tool') return;
+
+    const removeToolInAssistantMessage = async () => {
+      if (!message.parentId) return;
+      await get().internal_removeToolToAssistantMessage(message.parentId, message.tool_call_id);
+    };
+
+    await Promise.all([
+      // 1. remove tool message
+      get().internal_deleteMessage(id),
+      // 2. remove the tool item in the assistant tools
+      removeToolInAssistantMessage(),
+    ]);
+  },
+
   delAndRegenerateMessage: async (id) => {
     const traceId = chatSelectors.getTraceIdByMessageId(id)(get());
     get().internal_resendMessage(id, traceId);
@@ -197,7 +238,7 @@ export const chatMessage: StateCreator<
   clearMessage: async () => {
     const { activeId, activeTopicId, refreshMessages, refreshTopic, switchTopic } = get();
 
-    await messageService.removeMessages(activeId, activeTopicId);
+    await messageService.removeMessagesByAssistant(activeId, activeTopicId);
 
     if (activeTopicId) {
       await topicService.removeTopic(activeTopicId);
@@ -581,34 +622,7 @@ export const chatMessage: StateCreator<
       traceId: msgTraceId,
     };
   },
-  internal_toggleChatLoading: (loading, id, action) => {
-    return get().internal_toggleLoadingArrays('chatLoadingIds', loading, id, action);
-  },
-  internal_toggleMessageLoading: (loading, id) => {
-    set(
-      {
-        messageLoadingIds: toggleBooleanList(get().messageLoadingIds, id, loading),
-      },
-      false,
-      'internal_toggleMessageLoading',
-    );
-  },
-  internal_toggleToolCallingStreaming: (id, streaming) => {
-    set(
-      {
-        toolCallingStreamIds: produce(get().toolCallingStreamIds, (draft) => {
-          if (!!streaming) {
-            draft[id] = streaming;
-          } else {
-            delete draft[id];
-          }
-        }),
-      },
 
-      false,
-      'toggleToolCallingStreaming',
-    );
-  },
   internal_resendMessage: async (messageId, traceId) => {
     // 1. 构造所有相关的历史记录
     const chats = chatSelectors.currentChats(get());
@@ -715,7 +729,11 @@ export const chatMessage: StateCreator<
 
     return tempId;
   },
-
+  internal_deleteMessage: async (id: string) => {
+    get().internal_dispatchMessage({ type: 'deleteMessage', id });
+    await messageService.removeMessage(id);
+    await get().refreshMessages();
+  },
   internal_traceMessage: async (id, payload) => {
     // tracing the diff of update
     const message = chatSelectors.getMessageById(id)(get());
@@ -731,6 +749,35 @@ export const chatMessage: StateCreator<
     }
   },
 
+  // ----- Loading ------- //
+  internal_toggleMessageLoading: (loading, id) => {
+    set(
+      {
+        messageLoadingIds: toggleBooleanList(get().messageLoadingIds, id, loading),
+      },
+      false,
+      'internal_toggleMessageLoading',
+    );
+  },
+  internal_toggleChatLoading: (loading, id, action) => {
+    return get().internal_toggleLoadingArrays('chatLoadingIds', loading, id, action);
+  },
+  internal_toggleToolCallingStreaming: (id, streaming) => {
+    set(
+      {
+        toolCallingStreamIds: produce(get().toolCallingStreamIds, (draft) => {
+          if (!!streaming) {
+            draft[id] = streaming;
+          } else {
+            delete draft[id];
+          }
+        }),
+      },
+
+      false,
+      'toggleToolCallingStreaming',
+    );
+  },
   internal_toggleLoadingArrays: (key, loading, id, action) => {
     if (loading) {
       window.addEventListener('beforeunload', preventLeavingFn);

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

@@ -1,4 +1,4 @@
-import { ChatMessage } from '@/types/message';
+import { ChatMessage, ChatToolPayload } from '@/types/message';
 
 import { MessageDispatch, messagesReducer } from './reducer';
 
@@ -97,6 +97,35 @@ describe('messagesReducer', () => {
       expect(updatedMessage?.updatedAt).toBeGreaterThan(initialState[0].updatedAt);
     });
 
+    it('should update the extra field of a message if extra exist', () => {
+      const payload: MessageDispatch = {
+        type: 'updateMessageExtra',
+        id: 'data',
+        key: 'abc',
+        value: '2',
+      };
+
+      const newState = messagesReducer(
+        [
+          {
+            id: 'data',
+            content: 'Hello World',
+            createdAt: 1629264000000,
+            updatedAt: 1629264000000,
+            role: 'user',
+            meta: {},
+            extra: { abc: '1' },
+          } as ChatMessage,
+          ...initialState,
+        ],
+        payload,
+      );
+      const updatedMessage = newState.find((m) => m.id === 'data');
+
+      expect(updatedMessage?.extra).toEqual({ abc: '2' });
+      expect(updatedMessage?.updatedAt).toBeGreaterThan(initialState[0].updatedAt);
+    });
+
     it('should not modify state if message is not found', () => {
       const payload: MessageDispatch = {
         type: 'updateMessageExtra',
@@ -255,6 +284,151 @@ describe('messagesReducer', () => {
     });
   });
 
+  describe('addMessageTool', () => {
+    it('should add a tool to the specified assistant message if it dont have tools', () => {
+      const messageId = '1';
+      const toolPayload: ChatToolPayload = {
+        id: 'tc_1',
+        type: 'default',
+        identifier: 'tool1',
+        apiName: 'testFunction',
+        arguments: '{"arg1": "value1"}',
+      };
+
+      const payload: MessageDispatch = {
+        type: 'addMessageTool',
+        id: messageId,
+        value: toolPayload,
+      };
+
+      const newState = messagesReducer(
+        [...initialState, { id: messageId, role: 'assistant', content: '' } as ChatMessage],
+        payload,
+      );
+      const updatedMessage = newState.find((m) => m.id === messageId);
+
+      expect(updatedMessage).not.toBeUndefined();
+      expect(updatedMessage?.tools).toHaveLength(1);
+      expect(updatedMessage?.tools?.[0]).toEqual(toolPayload);
+      expect(updatedMessage?.updatedAt).toBeGreaterThan(initialState[0].updatedAt);
+    });
+
+    it('should add a tool to the specified assistant message', () => {
+      const messageId = 'message2';
+      const toolPayload: ChatToolPayload = {
+        id: 'tc_1',
+        type: 'default',
+        identifier: 'tool1',
+        apiName: 'testFunction',
+        arguments: '{"arg1": "value1"}',
+      };
+
+      const payload: MessageDispatch = {
+        type: 'addMessageTool',
+        id: messageId,
+        value: toolPayload,
+      };
+
+      const newState = messagesReducer(
+        [...initialState, { id: messageId, role: 'assistant', content: '' } as ChatMessage],
+        payload,
+      );
+      const updatedMessage = newState.find((m) => m.id === messageId);
+
+      expect(updatedMessage).not.toBeUndefined();
+      expect(updatedMessage?.tools).toHaveLength(2);
+      expect(updatedMessage?.tools?.[1]).toEqual(toolPayload);
+    });
+
+    it('should not modify the state if the message is not found', () => {
+      const toolPayload: ChatToolPayload = {
+        id: 'tc_1',
+        type: 'default',
+        identifier: 'tool1',
+        apiName: 'testFunction',
+        arguments: '{"arg1": "value1"}',
+      };
+
+      const payload: MessageDispatch = {
+        type: 'addMessageTool',
+        id: 'nonexistentMessage',
+        value: toolPayload,
+      };
+
+      const newState = messagesReducer(initialState, payload);
+      expect(newState).toEqual(initialState);
+    });
+
+    it('should not add a tool if the message is not an assistant message', () => {
+      const toolPayload: ChatToolPayload = {
+        id: 'tc_1',
+        type: 'default',
+        identifier: 'tool1',
+        apiName: 'testFunction',
+        arguments: '{"arg1": "value1"}',
+      };
+
+      const payload: MessageDispatch = {
+        type: 'addMessageTool',
+        id: 'message1', // This is a user message
+        value: toolPayload,
+      };
+
+      const newState = messagesReducer(initialState, payload);
+      expect(newState).toEqual(initialState);
+    });
+  });
+
+  describe('deleteMessageTool', () => {
+    it('should delete the specified tool from the message', () => {
+      const payload: MessageDispatch = {
+        type: 'deleteMessageTool',
+        id: 'message2',
+        tool_call_id: 'abc',
+      };
+
+      const newState = messagesReducer(initialState, payload);
+      const updatedMessage = newState.find((m) => m.id === 'message2');
+
+      expect(updatedMessage).not.toBeUndefined();
+      expect(updatedMessage?.tools).toHaveLength(0);
+      expect(updatedMessage?.updatedAt).toBeGreaterThan(initialState[0].updatedAt);
+    });
+
+    it('should not modify the state if the message is not found', () => {
+      const payload: MessageDispatch = {
+        type: 'deleteMessageTool',
+        id: 'nonexistentMessage',
+        tool_call_id: 'tool1',
+      };
+
+      const newState = messagesReducer(initialState, payload);
+      expect(newState).toEqual(initialState);
+    });
+
+    it('should not modify the state if the tool is not found', () => {
+      const payload: MessageDispatch = {
+        type: 'deleteMessageTool',
+        id: 'message1',
+        tool_call_id: 'nonexistentTool',
+      };
+
+      const newState = messagesReducer(initialState, payload);
+      expect(newState).toEqual(initialState);
+    });
+
+    it('should not delete a tool if the message is not an assistant message', () => {
+      const payload: MessageDispatch = {
+        type: 'deleteMessageTool',
+        id: 'message1', // This is a user message
+        tool_call_id: 'tool1',
+      };
+
+      const newState = messagesReducer(initialState, payload);
+      expect(newState).toEqual(initialState);
+    });
+  });
+
   describe('createMessage', () => {
     it('should add a new message to the state', () => {
       const payload: MessageDispatch = {
@@ -303,4 +477,29 @@ describe('messagesReducer', () => {
       expect(newState).toEqual(initialState);
     });
   });
+
+  describe('deleteMessages', () => {
+    it('should remove 2 messages from the state', () => {
+      const payload: MessageDispatch = {
+        type: 'deleteMessages',
+        ids: ['message1', 'message2'],
+      };
+
+      const newState = messagesReducer(initialState, payload);
+
+      expect(newState.length).toBe(0);
+      expect(newState.find((m) => m.id === 'message1')).toBeUndefined();
+      expect(newState.find((m) => m.id === 'message2')).toBeUndefined();
+    });
+
+    it('should not modify state if message to delete is not found', () => {
+      const payload: MessageDispatch = {
+        type: 'deleteMessage',
+        id: 'nonexistentMessage',
+      };
+
+      const newState = messagesReducer(initialState, payload);
+      expect(newState).toEqual(initialState);
+    });
+  });
 });