@lobehub/lobehub 2.0.0-next.35 → 2.0.0-next.37

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

@@ -0,0 +1,391 @@
+import { UIChatMessage } from '@lobechat/types';
+import { act, renderHook } from '@testing-library/react';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { chatService } from '@/services/chat';
+import { messageService } from '@/services/message';
+
+import { useChatStore } from '../../../../store';
+import { TEST_CONTENT, TEST_IDS, createMockMessage } from './fixtures';
+import { resetTestEnvironment, setupMockSelectors, spyOnMessageService } from './helpers';
+
+// Keep zustand mock as it's needed globally
+vi.mock('zustand/traditional');
+
+const realExecAgentRuntime = useChatStore.getState().internal_execAgentRuntime;
+
+beforeEach(() => {
+  resetTestEnvironment();
+  setupMockSelectors();
+  spyOnMessageService();
+
+  act(() => {
+    useChatStore.setState({
+      refreshMessages: vi.fn(),
+      internal_execAgentRuntime: vi.fn(),
+    });
+  });
+});
+
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+describe('StreamingExecutor actions', () => {
+  describe('internal_fetchAIChatMessage', () => {
+    it('should fetch and return AI chat response', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messages = [createMockMessage({ role: 'user' })];
+
+      const streamSpy = vi
+        .spyOn(chatService, 'createAssistantMessageStream')
+        .mockImplementation(async ({ onMessageHandle, onFinish }) => {
+          await onMessageHandle?.({ type: 'text', text: TEST_CONTENT.AI_RESPONSE } as any);
+          await onFinish?.(TEST_CONTENT.AI_RESPONSE, {});
+        });
+
+      await act(async () => {
+        const response = await result.current.internal_fetchAIChatMessage({
+          messages,
+          messageId: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          model: 'gpt-4o-mini',
+          provider: 'openai',
+        });
+        expect(response.isFunctionCall).toEqual(false);
+        expect(response.content).toEqual(TEST_CONTENT.AI_RESPONSE);
+      });
+
+      streamSpy.mockRestore();
+    });
+
+    it('should handle streaming errors gracefully', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messages = [createMockMessage({ role: 'user' })];
+
+      const streamSpy = vi
+        .spyOn(chatService, 'createAssistantMessageStream')
+        .mockImplementation(async ({ onErrorHandle }) => {
+          await onErrorHandle?.({ type: 'InvalidProviderAPIKey', message: 'Network error' } as any);
+        });
+
+      const updateMessageErrorSpy = vi.spyOn(messageService, 'updateMessageError');
+
+      await act(async () => {
+        await result.current.internal_fetchAIChatMessage({
+          model: 'gpt-4o-mini',
+          provider: 'openai',
+          messages,
+          messageId: TEST_IDS.ASSISTANT_MESSAGE_ID,
+        });
+      });
+
+      expect(updateMessageErrorSpy).toHaveBeenCalledWith(
+        TEST_IDS.ASSISTANT_MESSAGE_ID,
+        expect.objectContaining({ type: 'InvalidProviderAPIKey' }),
+      );
+
+      streamSpy.mockRestore();
+    });
+
+    it('should handle tool call chunks during streaming', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messages = [createMockMessage({ role: 'user' })];
+
+      const streamSpy = vi
+        .spyOn(chatService, 'createAssistantMessageStream')
+        .mockImplementation(async ({ onMessageHandle, onFinish }) => {
+          await onMessageHandle?.({
+            type: 'tool_calls',
+            isAnimationActives: [true],
+            tool_calls: [
+              { id: 'tool-1', type: 'function', function: { name: 'test', arguments: '{}' } },
+            ],
+          } as any);
+          await onFinish?.(TEST_CONTENT.AI_RESPONSE, {
+            toolCalls: [
+              { id: 'tool-1', type: 'function', function: { name: 'test', arguments: '{}' } },
+            ],
+          } as any);
+        });
+
+      await act(async () => {
+        const response = await result.current.internal_fetchAIChatMessage({
+          messages,
+          messageId: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          model: 'gpt-4o-mini',
+          provider: 'openai',
+        });
+        expect(response.isFunctionCall).toEqual(true);
+      });
+
+      streamSpy.mockRestore();
+    });
+
+    it('should handle text chunks during streaming', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messages = [createMockMessage({ role: 'user' })];
+      const dispatchSpy = vi.spyOn(result.current, 'internal_dispatchMessage');
+
+      const streamSpy = vi
+        .spyOn(chatService, 'createAssistantMessageStream')
+        .mockImplementation(async ({ onMessageHandle, onFinish }) => {
+          await onMessageHandle?.({ type: 'text', text: 'Hello' } as any);
+          await onMessageHandle?.({ type: 'text', text: ' World' } as any);
+          await onFinish?.('Hello World', {} as any);
+        });
+
+      await act(async () => {
+        await result.current.internal_fetchAIChatMessage({
+          messages,
+          messageId: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          model: 'gpt-4o-mini',
+          provider: 'openai',
+        });
+      });
+
+      expect(dispatchSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          id: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          type: 'updateMessage',
+          value: expect.objectContaining({ content: 'Hello' }),
+        }),
+      );
+
+      streamSpy.mockRestore();
+    });
+
+    it('should handle reasoning chunks during streaming', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messages = [createMockMessage({ role: 'user' })];
+      const dispatchSpy = vi.spyOn(result.current, 'internal_dispatchMessage');
+
+      const streamSpy = vi
+        .spyOn(chatService, 'createAssistantMessageStream')
+        .mockImplementation(async ({ onMessageHandle, onFinish }) => {
+          await onMessageHandle?.({ type: 'reasoning', text: 'Thinking...' } as any);
+          await onMessageHandle?.({ type: 'text', text: 'Answer' } as any);
+          await onFinish?.('Answer', {} as any);
+        });
+
+      await act(async () => {
+        await result.current.internal_fetchAIChatMessage({
+          messages,
+          messageId: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          model: 'gpt-4o-mini',
+          provider: 'openai',
+        });
+      });
+
+      expect(dispatchSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          id: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          type: 'updateMessage',
+          value: expect.objectContaining({ reasoning: { content: 'Thinking...' } }),
+        }),
+      );
+
+      streamSpy.mockRestore();
+    });
+
+    it('should skip grounding when citations are empty', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messages = [createMockMessage({ role: 'user' })];
+      const dispatchSpy = vi.spyOn(result.current, 'internal_dispatchMessage');
+
+      const streamSpy = vi
+        .spyOn(chatService, 'createAssistantMessageStream')
+        .mockImplementation(async ({ onMessageHandle, onFinish }) => {
+          await onMessageHandle?.({
+            type: 'grounding',
+            grounding: { citations: [], searchQueries: [] },
+          } as any);
+          await onFinish?.('Answer', {} as any);
+        });
+
+      await act(async () => {
+        await result.current.internal_fetchAIChatMessage({
+          messages,
+          messageId: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          model: 'gpt-4o-mini',
+          provider: 'openai',
+        });
+      });
+
+      // Should not dispatch when citations are empty
+      const groundingCalls = dispatchSpy.mock.calls.filter((call) => {
+        const dispatch = call[0];
+        return dispatch?.type === 'updateMessage' && 'value' in dispatch && dispatch.value?.search;
+      });
+      expect(groundingCalls).toHaveLength(0);
+
+      streamSpy.mockRestore();
+    });
+
+    it('should handle grounding chunks during streaming', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messages = [createMockMessage({ role: 'user' })];
+      const dispatchSpy = vi.spyOn(result.current, 'internal_dispatchMessage');
+
+      const streamSpy = vi
+        .spyOn(chatService, 'createAssistantMessageStream')
+        .mockImplementation(async ({ onMessageHandle, onFinish }) => {
+          await onMessageHandle?.({
+            type: 'grounding',
+            grounding: {
+              citations: [{ url: 'https://example.com', title: 'Example' }],
+              searchQueries: ['test query'],
+            },
+          } as any);
+          await onFinish?.('Answer', {} as any);
+        });
+
+      await act(async () => {
+        await result.current.internal_fetchAIChatMessage({
+          messages,
+          messageId: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          model: 'gpt-4o-mini',
+          provider: 'openai',
+        });
+      });
+
+      expect(dispatchSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          id: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          type: 'updateMessage',
+          value: expect.objectContaining({
+            search: expect.objectContaining({
+              citations: expect.any(Array),
+            }),
+          }),
+        }),
+      );
+
+      streamSpy.mockRestore();
+    });
+
+    it('should handle base64 image chunks during streaming', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messages = [createMockMessage({ role: 'user' })];
+      const dispatchSpy = vi.spyOn(result.current, 'internal_dispatchMessage');
+
+      const streamSpy = vi
+        .spyOn(chatService, 'createAssistantMessageStream')
+        .mockImplementation(async ({ onMessageHandle, onFinish }) => {
+          await onMessageHandle?.({
+            type: 'base64_image',
+            image: { id: 'img-1', data: 'base64data' },
+            images: [{ id: 'img-1', data: 'base64data' }],
+          } as any);
+          await onFinish?.('Answer', {} as any);
+        });
+
+      await act(async () => {
+        await result.current.internal_fetchAIChatMessage({
+          messages,
+          messageId: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          model: 'gpt-4o-mini',
+          provider: 'openai',
+        });
+      });
+
+      expect(dispatchSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          id: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          type: 'updateMessage',
+          value: expect.objectContaining({
+            imageList: expect.any(Array),
+          }),
+        }),
+      );
+
+      streamSpy.mockRestore();
+    });
+
+    it('should handle empty tool call arguments', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messages = [createMockMessage({ role: 'user' })];
+
+      const streamSpy = vi
+        .spyOn(chatService, 'createAssistantMessageStream')
+        .mockImplementation(async ({ onFinish }) => {
+          await onFinish?.(TEST_CONTENT.AI_RESPONSE, {
+            toolCalls: [
+              { id: 'tool-1', type: 'function', function: { name: 'test', arguments: '' } },
+            ],
+          } as any);
+        });
+
+      await act(async () => {
+        const response = await result.current.internal_fetchAIChatMessage({
+          messages,
+          messageId: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          model: 'gpt-4o-mini',
+          provider: 'openai',
+        });
+        expect(response.isFunctionCall).toEqual(true);
+      });
+
+      streamSpy.mockRestore();
+    });
+
+    it('should update message with traceId when provided in onFinish', async () => {
+      const { result } = renderHook(() => useChatStore());
+      const messages = [createMockMessage({ role: 'user' })];
+      const traceId = 'test-trace-123';
+
+      const updateMessageSpy = vi.spyOn(messageService, 'updateMessage');
+      const streamSpy = vi
+        .spyOn(chatService, 'createAssistantMessageStream')
+        .mockImplementation(async ({ onFinish }) => {
+          await onFinish?.(TEST_CONTENT.AI_RESPONSE, { traceId } as any);
+        });
+
+      await act(async () => {
+        await result.current.internal_fetchAIChatMessage({
+          messages,
+          messageId: TEST_IDS.ASSISTANT_MESSAGE_ID,
+          model: 'gpt-4o-mini',
+          provider: 'openai',
+        });
+      });
+
+      expect(updateMessageSpy).toHaveBeenCalledWith(
+        TEST_IDS.ASSISTANT_MESSAGE_ID,
+        expect.objectContaining({ traceId }),
+      );
+
+      streamSpy.mockRestore();
+    });
+  });
+
+  describe('internal_execAgentRuntime', () => {
+    it('should handle the core AI message processing', async () => {
+      act(() => {
+        useChatStore.setState({ internal_execAgentRuntime: realExecAgentRuntime });
+      });
+
+      const { result } = renderHook(() => useChatStore());
+      const userMessage = {
+        id: TEST_IDS.USER_MESSAGE_ID,
+        role: 'user',
+        content: TEST_CONTENT.USER_MESSAGE,
+        sessionId: TEST_IDS.SESSION_ID,
+        topicId: TEST_IDS.TOPIC_ID,
+      } as UIChatMessage;
+      const messages = [userMessage];
+
+      const streamSpy = vi.spyOn(chatService, 'createAssistantMessageStream');
+
+      await act(async () => {
+        await result.current.internal_execAgentRuntime({
+          messages,
+          parentMessageId: userMessage.id,
+          parentMessageType: 'user',
+        });
+      });
+
+      expect(streamSpy).toHaveBeenCalled();
+      expect(result.current.refreshMessages).toHaveBeenCalled();
+    });
+  });
+});

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

@@ -0,0 +1,179 @@
+import { act, renderHook } from '@testing-library/react';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { useChatStore } from '../../../../store';
+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('StreamingStates actions', () => {
+  describe('internal_toggleChatLoading', () => {
+    it('should enable loading state with new abort controller', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.internal_toggleChatLoading(true, TEST_IDS.MESSAGE_ID, 'test-action');
+      });
+
+      const state = useChatStore.getState();
+      expect(state.chatLoadingIdsAbortController).toBeInstanceOf(AbortController);
+      expect(state.chatLoadingIds).toEqual([TEST_IDS.MESSAGE_ID]);
+    });
+
+    it('should disable loading state and clear abort controller', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.internal_toggleChatLoading(true, TEST_IDS.MESSAGE_ID, 'start');
+        result.current.internal_toggleChatLoading(false, undefined, 'stop');
+      });
+
+      const state = useChatStore.getState();
+      expect(state.chatLoadingIdsAbortController).toBeUndefined();
+      expect(state.chatLoadingIds).toEqual([]);
+    });
+
+    it('should manage beforeunload event listener', () => {
+      const { result } = renderHook(() => useChatStore());
+      const addListenerSpy = vi.spyOn(window, 'addEventListener');
+      const removeListenerSpy = vi.spyOn(window, 'removeEventListener');
+
+      act(() => {
+        result.current.internal_toggleChatLoading(true, TEST_IDS.MESSAGE_ID, 'start');
+      });
+
+      expect(addListenerSpy).toHaveBeenCalledWith('beforeunload', expect.any(Function));
+
+      act(() => {
+        result.current.internal_toggleChatLoading(false, undefined, 'stop');
+      });
+
+      expect(removeListenerSpy).toHaveBeenCalledWith('beforeunload', expect.any(Function));
+    });
+
+    it('should reuse existing abort controller', () => {
+      const existingController = new AbortController();
+
+      act(() => {
+        useChatStore.setState({ chatLoadingIdsAbortController: existingController });
+      });
+
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.internal_toggleChatLoading(true, TEST_IDS.MESSAGE_ID, 'test');
+      });
+
+      const state = useChatStore.getState();
+      expect(state.chatLoadingIdsAbortController).toStrictEqual(existingController);
+    });
+  });
+
+  describe('internal_toggleToolCallingStreaming', () => {
+    it('should track tool calling stream status', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.internal_toggleToolCallingStreaming(TEST_IDS.MESSAGE_ID, [true]);
+      });
+
+      expect(result.current.toolCallingStreamIds[TEST_IDS.MESSAGE_ID]).toEqual([true]);
+    });
+
+    it('should clear tool calling stream status', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.internal_toggleToolCallingStreaming(TEST_IDS.MESSAGE_ID, [true]);
+        result.current.internal_toggleToolCallingStreaming(TEST_IDS.MESSAGE_ID, undefined);
+      });
+
+      expect(result.current.toolCallingStreamIds[TEST_IDS.MESSAGE_ID]).toBeUndefined();
+    });
+  });
+
+  describe('internal_toggleSearchWorkflow', () => {
+    it('should enable search workflow loading state', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.internal_toggleSearchWorkflow(true, TEST_IDS.MESSAGE_ID);
+      });
+
+      const state = useChatStore.getState();
+      expect(state.searchWorkflowLoadingIds).toEqual([TEST_IDS.MESSAGE_ID]);
+    });
+
+    it('should disable search workflow loading state', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.internal_toggleSearchWorkflow(true, TEST_IDS.MESSAGE_ID);
+        result.current.internal_toggleSearchWorkflow(false, TEST_IDS.MESSAGE_ID);
+      });
+
+      const state = useChatStore.getState();
+      expect(state.searchWorkflowLoadingIds).toEqual([]);
+    });
+  });
+
+  describe('internal_toggleChatReasoning', () => {
+    it('should enable reasoning loading state', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.internal_toggleChatReasoning(true, TEST_IDS.MESSAGE_ID, 'test-action');
+      });
+
+      const state = useChatStore.getState();
+      expect(state.reasoningLoadingIds).toEqual([TEST_IDS.MESSAGE_ID]);
+    });
+
+    it('should disable reasoning loading state', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.internal_toggleChatReasoning(true, TEST_IDS.MESSAGE_ID, 'start');
+        result.current.internal_toggleChatReasoning(false, TEST_IDS.MESSAGE_ID, 'stop');
+      });
+
+      const state = useChatStore.getState();
+      expect(state.reasoningLoadingIds).toEqual([]);
+    });
+  });
+
+  describe('internal_toggleMessageInToolsCalling', () => {
+    it('should enable tools calling state', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.internal_toggleMessageInToolsCalling(true, TEST_IDS.MESSAGE_ID);
+      });
+
+      const state = useChatStore.getState();
+      expect(state.messageInToolsCallingIds).toEqual([TEST_IDS.MESSAGE_ID]);
+    });
+
+    it('should disable tools calling state', () => {
+      const { result } = renderHook(() => useChatStore());
+
+      act(() => {
+        result.current.internal_toggleMessageInToolsCalling(true, TEST_IDS.MESSAGE_ID);
+        result.current.internal_toggleMessageInToolsCalling(false, TEST_IDS.MESSAGE_ID);
+      });
+
+      const state = useChatStore.getState();
+      expect(state.messageInToolsCallingIds).toEqual([]);
+    });
+  });
+});