From ad4fe8407838309fc24d92837be058fabc1c8b3b Mon Sep 17 00:00:00 2001 From: Steph Milovic Date: Tue, 2 Jul 2024 13:42:22 -0600 Subject: [PATCH] [Security solution] Assistant race condition bug fixing (#187186) --- .../assistant_header_flyout.tsx | 3 ++ .../assistant/chat_send/use_chat_send.tsx | 2 +- .../impl/assistant/index.test.tsx | 37 ++++++++++++------- .../impl/assistant/index.tsx | 25 ++++++------- 4 files changed, 39 insertions(+), 28 deletions(-) diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_header/assistant_header_flyout.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_header/assistant_header_flyout.tsx index d9ba27b96655f..bd75e80aef0ca 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_header/assistant_header_flyout.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_header/assistant_header_flyout.tsx @@ -125,6 +125,7 @@ export const AssistantHeaderFlyout: React.FC = ({ `, onClick: showDestroyModal, icon: 'refresh', + 'data-test-subj': 'clear-chat', }, ], }, @@ -243,6 +244,7 @@ export const AssistantHeaderFlyout: React.FC = ({ aria-label="test" iconType="boxesVertical" onClick={onButtonClick} + data-test-subj="chat-context-menu" /> } isOpen={isPopoverOpen} @@ -266,6 +268,7 @@ export const AssistantHeaderFlyout: React.FC = ({ confirmButtonText={i18n.RESET_BUTTON_TEXT} buttonColor="danger" defaultFocusedButton="confirm" + data-test-subj="reset-conversation-modal" >

{i18n.CLEAR_CHAT_CONFIRMATION}

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 index 020822821d163..f42fe17242d86 100644 --- 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; 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 index c023970803da4..b25945dd247bf 100644 --- 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); 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, 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 index 830c5d2b7080a..907e4d70accd5 100644 --- 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 = ({ ); 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 = ({ }); } }, [ + areConnectorsFetched, conversationTitle, conversations, - getDefaultConversation, - getLastConversationId, conversationsLoaded, - currentConversation?.id, currentConversationId, + getDefaultConversation, + getLastConversationId, isAssistantEnabled, isFlyoutMode, ]); @@ -549,7 +549,7 @@ const AssistantComponent: React.FC = ({ const { abortStream, - handleOnChatCleared, + handleOnChatCleared: onChatCleared, handlePromptChange, handleSendMessage, handleRegenerateResponse, @@ -567,6 +567,11 @@ const AssistantComponent: React.FC = ({ 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 = ({ } } })(); - }, [ - currentConversation, - defaultConnector, - refetchConversationsState, - setApiConfig, - showMissingConnectorCallout, - areConnectorsFetched, - mutateAsync, - ]); + }, [areConnectorsFetched, currentConversation, mutateAsync]); const handleCreateConversation = useCallback(async () => { const newChatExists = find(conversations, ['title', NEW_CHAT]);