[Security solution] Assistant race condition bug fixing (#187186)

elastic · Jul 2, 2024 · ad4fe84 · ad4fe84
1 parent 2aa94a2
commit ad4fe84
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 28 deletions.
diff --git a/...ackages/kbn-elastic-assistant/impl/assistant/assistant_header/assistant_header_flyout.tsx b/...ackages/kbn-elastic-assistant/impl/assistant/assistant_header/assistant_header_flyout.tsx
@@ -125,6 +125,7 @@ export const AssistantHeaderFlyout: React.FC<Props> = ({
             `,
             onClick: showDestroyModal,
             icon: 'refresh',
+            'data-test-subj': 'clear-chat',
           },
         ],
       },
@@ -243,6 +244,7 @@ export const AssistantHeaderFlyout: React.FC<Props> = ({
                       aria-label="test"
                       iconType="boxesVertical"
                       onClick={onButtonClick}
+                      data-test-subj="chat-context-menu"
                     />
                   }
                   isOpen={isPopoverOpen}
@@ -266,6 +268,7 @@ export const AssistantHeaderFlyout: React.FC<Props> = ({
           confirmButtonText={i18n.RESET_BUTTON_TEXT}
           buttonColor="danger"
           defaultFocusedButton="confirm"
+          data-test-subj="reset-conversation-modal"
         >
           <p>{i18n.CLEAR_CHAT_CONFIRMATION}</p>
         </EuiConfirmModal>

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/chat_send/use_chat_send.tsx
@@ -35,7 +35,7 @@ export interface UseChatSendProps {
 
 export interface UseChatSend {
   abortStream: () => void;
-  handleOnChatCleared: () => void;
+  handleOnChatCleared: () => Promise<void>;
   handlePromptChange: (prompt: string) => void;
   handleSendMessage: (promptText: string) => void;
   handleRegenerateResponse: () => void;

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx
@@ -7,7 +7,7 @@
 
 import React from 'react';
 
-import { act, fireEvent, render, screen, within } from '@testing-library/react';
+import { act, fireEvent, render, screen, waitFor, within } from '@testing-library/react';
 import { Assistant } from '.';
 import type { IHttpFetchError } from '@kbn/core/public';
 
@@ -63,15 +63,25 @@ const mockData = {
   },
 };
 const mockDeleteConvo = jest.fn();
+const clearConversation = jest.fn();
 const mockUseConversation = {
+  clearConversation: clearConversation.mockResolvedValue(mockData.welcome_id),
   getConversation: jest.fn(),
   getDefaultConversation: jest.fn().mockReturnValue(mockData.welcome_id),
   deleteConversation: mockDeleteConvo,
   setApiConfig: jest.fn().mockResolvedValue({}),
 };
 
+const refetchResults = jest.fn();
+
 describe('Assistant', () => {
-  beforeAll(() => {
+  let persistToLocalStorage: jest.Mock;
+  let persistToSessionStorage: jest.Mock;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    persistToLocalStorage = jest.fn();
+    persistToSessionStorage = jest.fn();
     (useConversation as jest.Mock).mockReturnValue(mockUseConversation);
     jest.mocked(useConnectorSetup).mockReturnValue({
       comments: [],
@@ -89,13 +99,14 @@ describe('Assistant', () => {
     ];
     jest.mocked(useLoadConnectors).mockReturnValue({
       isFetched: true,
+      isFetchedAfterMount: true,
       data: connectors,
     } as unknown as UseQueryResult<AIConnector[], IHttpFetchError>);
 
     jest.mocked(useFetchCurrentUserConversations).mockReturnValue({
       data: mockData,
       isLoading: false,
-      refetch: jest.fn().mockResolvedValue({
+      refetch: refetchResults.mockResolvedValue({
         isLoading: false,
         data: {
           ...mockData,
@@ -107,16 +118,6 @@ describe('Assistant', () => {
       }),
       isFetched: true,
     } as unknown as DefinedUseQueryResult<Record<string, Conversation>, unknown>);
-  });
-
-  let persistToLocalStorage: jest.Mock;
-  let persistToSessionStorage: jest.Mock;
-
-  beforeEach(() => {
-    jest.clearAllMocks();
-    persistToLocalStorage = jest.fn();
-    persistToSessionStorage = jest.fn();
-
     jest
       .mocked(useLocalStorage)
       .mockReturnValue([undefined, persistToLocalStorage] as unknown as ReturnType<
@@ -234,6 +235,16 @@ describe('Assistant', () => {
       });
       expect(mockDeleteConvo).toHaveBeenCalledWith(mockData.welcome_id.id);
     });
+    it('should refetchConversationsState after clear chat history button click', async () => {
+      renderAssistant({ isFlyoutMode: true });
+      fireEvent.click(screen.getByTestId('chat-context-menu'));
+      fireEvent.click(screen.getByTestId('clear-chat'));
+      fireEvent.click(screen.getByTestId('confirmModalConfirmButton'));
+      await waitFor(() => {
+        expect(clearConversation).toHaveBeenCalled();
+        expect(refetchResults).toHaveBeenCalled();
+      });
+    });
   });
   describe('when selected conversation changes and some connectors are loaded', () => {
     it('should persist the conversation id to local storage', async () => {

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx
@@ -220,7 +220,7 @@ const AssistantComponent: React.FC<Props> = ({
   );
 
   useEffect(() => {
-    if (conversationsLoaded && Object.keys(conversations).length > 0) {
+    if (areConnectorsFetched && conversationsLoaded && Object.keys(conversations).length > 0) {
       setCurrentConversation((prev) => {
         const nextConversation =
           (currentConversationId && conversations[currentConversationId]) ||
@@ -256,13 +256,13 @@ const AssistantComponent: React.FC<Props> = ({
       });
     }
   }, [
+    areConnectorsFetched,
     conversationTitle,
     conversations,
-    getDefaultConversation,
-    getLastConversationId,
     conversationsLoaded,
-    currentConversation?.id,
     currentConversationId,
+    getDefaultConversation,
+    getLastConversationId,
     isAssistantEnabled,
     isFlyoutMode,
   ]);
@@ -549,7 +549,7 @@ const AssistantComponent: React.FC<Props> = ({
 
   const {
     abortStream,
-    handleOnChatCleared,
+    handleOnChatCleared: onChatCleared,
     handlePromptChange,
     handleSendMessage,
     handleRegenerateResponse,
@@ -567,6 +567,11 @@ const AssistantComponent: React.FC<Props> = ({
     setCurrentConversation,
   });
 
+  const handleOnChatCleared = useCallback(async () => {
+    await onChatCleared();
+    await refetchResults();
+  }, [onChatCleared, refetchResults]);
+
   const handleChatSend = useCallback(
     async (promptText: string) => {
       await handleSendMessage(promptText);
@@ -733,15 +738,7 @@ const AssistantComponent: React.FC<Props> = ({
         }
       }
     })();
-  }, [
-    currentConversation,
-    defaultConnector,
-    refetchConversationsState,
-    setApiConfig,
-    showMissingConnectorCallout,
-    areConnectorsFetched,
-    mutateAsync,
-  ]);
+  }, [areConnectorsFetched, currentConversation, mutateAsync]);
 
   const handleCreateConversation = useCallback(async () => {
     const newChatExists = find(conversations, ['title', NEW_CHAT]);