[Security Solution] Store last conversation in localstorage #6993

x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_overlay/index.tsx

@@ -32,7 +32,7 @@ export const AssistantOverlay: React.FC = React.memo(() => {
     WELCOME_CONVERSATION_TITLE
   );
   const [promptContextId, setPromptContextId] = useState<string | undefined>();
-  const { setShowAssistantOverlay } = useAssistantContext();
+  const { setShowAssistantOverlay, localStorageLastConversationId } = useAssistantContext();
 
   // Bind `showAssistantOverlay` in SecurityAssistantContext to this modal instance
   const showOverlay = useCallback(
@@ -52,15 +52,25 @@ export const AssistantOverlay: React.FC = React.memo(() => {
     setShowAssistantOverlay(showOverlay);
   }, [setShowAssistantOverlay, showOverlay]);
 
+  // Called whenever platform specific shortcut for assistant is pressed
+  const handleShortcutPress = useCallback(() => {
+    // Try to restore the last conversation on shortcut pressed
+    if (!isModalVisible) {
+      setConversationId(localStorageLastConversationId || WELCOME_CONVERSATION_TITLE);
+    }
+
+    setIsModalVisible(!isModalVisible);
+  }, [isModalVisible, localStorageLastConversationId]);
+
   // Register keyboard listener to show the modal when cmd + ; is pressed
   const onKeyDown = useCallback(
     (event: KeyboardEvent) => {
       if (event.key === ';' && (isMac ? event.metaKey : event.ctrlKey)) {
         event.preventDefault();
-        setIsModalVisible(!isModalVisible);
+        handleShortcutPress();
       }
     },
-    [isModalVisible]
+    [handleShortcutPress]
   );
   useEvent('keydown', onKeyDown);
 

x-pack/packages/kbn-elastic-assistant/impl/assistant/index.test.tsx

@@ -0,0 +1,123 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import React from 'react';
+
+import { act, fireEvent, render, screen } from '@testing-library/react';
+import { Assistant } from '.';
+import { Conversation } from '../assistant_context/types';
+import { IHttpFetchError } from '@kbn/core/public';
+import { ActionConnector } from '@kbn/triggers-actions-ui-plugin/public';
+
+import { useLoadConnectors } from '../connectorland/use_load_connectors';
+import { useConnectorSetup } from '../connectorland/connector_setup';
+
+import { UseQueryResult } from '@tanstack/react-query';
+import { WELCOME_CONVERSATION_TITLE } from './use_conversation/translations';
+
+import { useLocalStorage } from 'react-use';
+import { PromptEditor } from './prompt_editor';
+import { QuickPrompts } from './quick_prompts/quick_prompts';
+import { TestProviders } from '../mock/test_providers/test_providers';
+
+jest.mock('../connectorland/use_load_connectors');
+jest.mock('../connectorland/connector_setup');
+jest.mock('react-use');
+
+jest.mock('./prompt_editor', () => ({ PromptEditor: jest.fn() }));
+jest.mock('./quick_prompts/quick_prompts', () => ({ QuickPrompts: jest.fn() }));
+
+const MOCK_CONVERSATION_TITLE = 'electric sheep';
+
+const getInitialConversations = (): Record<string, Conversation> => ({
+  [WELCOME_CONVERSATION_TITLE]: {
+    id: WELCOME_CONVERSATION_TITLE,
+    messages: [],
+    apiConfig: {},
+  },
+  [MOCK_CONVERSATION_TITLE]: {
+    id: MOCK_CONVERSATION_TITLE,
+    messages: [],
+    apiConfig: {},
+  },
+});
+
+const renderAssistant = () =>
+  render(
+    <TestProviders getInitialConversations={getInitialConversations}>
+      <Assistant />
+    </TestProviders>
+  );
+
+describe('Assistant', () => {
+  beforeAll(() => {
+    jest.mocked(useConnectorSetup).mockReturnValue({
+      connectorDialog: <></>,
+      connectorPrompt: <></>,
+    });
+
+    jest.mocked(PromptEditor).mockReturnValue(null);
+    jest.mocked(QuickPrompts).mockReturnValue(null);
+  });
+
+  beforeEach(jest.clearAllMocks);
+
+  let persistToLocalStorage: jest.Mock;
+
+  beforeEach(() => {
+    persistToLocalStorage = jest.fn();
+
+    jest
+      .mocked(useLocalStorage)
+      .mockReturnValue(['', persistToLocalStorage] as unknown as ReturnType<
+        typeof useLocalStorage
+      >);
+  });
+
+  describe('when selected conversation changes and some connectors are loaded', () => {
+    it('should persist the conversation id to local storage', async () => {
+      const connectors: unknown[] = [{}];
+
+      jest.mocked(useLoadConnectors).mockReturnValue({
+        isSuccess: true,
+        data: connectors,
+      } as unknown as UseQueryResult<ActionConnector[], IHttpFetchError>);
+
+      renderAssistant();
+
+      expect(persistToLocalStorage).toHaveBeenCalled();
+
+      expect(persistToLocalStorage).toHaveBeenLastCalledWith(WELCOME_CONVERSATION_TITLE);
+
+      const previousConversationButton = screen.getByLabelText('Previous conversation');
+
+      expect(previousConversationButton).toBeInTheDocument();
+
+      await act(async () => {
+        fireEvent.click(previousConversationButton);
+      });
+
+      expect(persistToLocalStorage).toHaveBeenLastCalledWith('electric sheep');
+    });
+  });
+
+  describe('when no connectors are loaded', () => {
+    it('should clear conversation id in local storage', async () => {
+      const emptyConnectors: unknown[] = [];
+
+      jest.mocked(useLoadConnectors).mockReturnValue({
+        isSuccess: true,
+        data: emptyConnectors,
+      } as unknown as UseQueryResult<ActionConnector[], IHttpFetchError>);
+
+      renderAssistant();
+
+      expect(persistToLocalStorage).toHaveBeenCalled();
+      expect(persistToLocalStorage).toHaveBeenLastCalledWith('');
+    });
+  });
+});

x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx

@@ -78,6 +78,7 @@ const AssistantComponent: React.FC<Props> = ({
     getComments,
     http,
     promptContexts,
+    setLastConversationId,
     title,
   } = useAssistantContext();
   const [selectedPromptContexts, setSelectedPromptContexts] = useState<
@@ -105,7 +106,11 @@ const AssistantComponent: React.FC<Props> = ({
     [conversations, selectedConversationId]
   );
 
-  const { data: connectors, refetch: refetchConnectors } = useLoadConnectors({ http });
+  const {
+    data: connectors,
+    isSuccess: areConnectorsFetched,
+    refetch: refetchConnectors,
+  } = useLoadConnectors({ http });
   const defaultConnectorId = useMemo(() => connectors?.[0]?.id, [connectors]);
   const defaultProvider = useMemo(
     () =>
@@ -114,6 +119,15 @@ const AssistantComponent: React.FC<Props> = ({
     [connectors]
   );
 
+  // Remember last selection for reuse after keyboard shortcut is pressed.
+  // Clear it if there is no connectors
+  useEffect(() => {
+    if (areConnectorsFetched && !connectors?.length) {
+      return setLastConversationId('');
+    }
+    setLastConversationId(selectedConversationId);
+  }, [areConnectorsFetched, connectors?.length, selectedConversationId, setLastConversationId]);
+
   const isWelcomeSetup = (connectors?.length ?? 0) === 0;
 
   const { connectorDialog, connectorPrompt } = useConnectorSetup({
@@ -178,11 +192,11 @@ const AssistantComponent: React.FC<Props> = ({
 
   // Scroll to bottom on conversation change
   useEffect(() => {
-    bottomRef.current?.scrollIntoView({ behavior: 'auto' });
+    bottomRef.current?.scrollIntoView?.({ behavior: 'auto' });
   }, []);
   useEffect(() => {
     setTimeout(() => {
-      bottomRef.current?.scrollIntoView({ behavior: 'auto' });
+      bottomRef.current?.scrollIntoView?.({ behavior: 'auto' });
       promptTextAreaRef?.current?.focus();
     }, 0);
   }, [currentConversation.messages.length, selectedPromptContextsCount]);

x-pack/packages/kbn-elastic-assistant/impl/assistant_context/constants.tsx

@@ -8,3 +8,4 @@
 export const DEFAULT_ASSISTANT_NAMESPACE = 'elasticAssistantDefault';
 export const QUICK_PROMPT_LOCAL_STORAGE_KEY = 'quickPrompts';
 export const SYSTEM_PROMPT_LOCAL_STORAGE_KEY = 'systemPrompts';
+export const LAST_CONVERSATION_ID_LOCAL_STORAGE_KEY = 'lastConversationId';

x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx

@@ -27,6 +27,7 @@ import { Prompt } from '../assistant/types';
 import { BASE_SYSTEM_PROMPTS } from '../content/prompts/system';
 import {
   DEFAULT_ASSISTANT_NAMESPACE,
+  LAST_CONVERSATION_ID_LOCAL_STORAGE_KEY,
   QUICK_PROMPT_LOCAL_STORAGE_KEY,
   SYSTEM_PROMPT_LOCAL_STORAGE_KEY,
 } from './constants';
@@ -96,6 +97,7 @@ interface UseAssistantContext {
     showAnonymizedValues: boolean;
   }) => EuiCommentProps[];
   http: HttpSetup;
+  localStorageLastConversationId: string | undefined;
   promptContexts: Record<string, PromptContext>;
   nameSpace: string;
   registerPromptContext: RegisterPromptContext;
@@ -104,6 +106,7 @@ interface UseAssistantContext {
   setConversations: React.Dispatch<React.SetStateAction<Record<string, Conversation>>>;
   setDefaultAllow: React.Dispatch<React.SetStateAction<string[]>>;
   setDefaultAllowReplacement: React.Dispatch<React.SetStateAction<string[]>>;
+  setLastConversationId: React.Dispatch<React.SetStateAction<string | undefined>>;
   setShowAssistantOverlay: (showAssistantOverlay: ShowAssistantOverlay) => void;
   showAssistantOverlay: ShowAssistantOverlay;
   title: string;
@@ -148,6 +151,9 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
     baseSystemPrompts
   );
 
+  const [localStorageLastConversationId, setLocalStorageLastConversationId] =
+    useLocalStorage<string>(`${nameSpace}.${LAST_CONVERSATION_ID_LOCAL_STORAGE_KEY}`);
+
   /**
    * Prompt contexts are used to provide components a way to register and make their data available to the assistant.
    */
@@ -248,6 +254,8 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       showAssistantOverlay,
       title,
       unRegisterPromptContext,
+      localStorageLastConversationId,
+      setLastConversationId: setLocalStorageLastConversationId,
     }),
     [
       actionTypeRegistry,
@@ -263,6 +271,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       defaultAllowReplacement,
       getComments,
       http,
+      localStorageLastConversationId,
       localStorageQuickPrompts,
       localStorageSystemPrompts,
       nameSpace,
@@ -271,6 +280,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       registerPromptContext,
       setDefaultAllow,
       setDefaultAllowReplacement,
+      setLocalStorageLastConversationId,
       setLocalStorageQuickPrompts,
       setLocalStorageSystemPrompts,
       showAssistantOverlay,

x-pack/packages/kbn-elastic-assistant/impl/mock/test_providers/test_providers.tsx

@@ -14,17 +14,23 @@ import React from 'react';
 import { ThemeProvider } from 'styled-components';
 
 import { AssistantProvider } from '../../assistant_context';
+import { Conversation } from '../../assistant_context/types';
 
 interface Props {
   children: React.ReactNode;
+  getInitialConversations?: () => Record<string, Conversation>;
 }
 
 window.scrollTo = jest.fn();
 
+const mockGetInitialConversations = () => ({});
+
 /** A utility for wrapping children in the providers required to run tests */
-export const TestProvidersComponent: React.FC<Props> = ({ children }) => {
+export const TestProvidersComponent: React.FC<Props> = ({
+  children,
+  getInitialConversations = mockGetInitialConversations,
+}) => {
   const actionTypeRegistry = actionTypeRegistryMock.create();
-  const mockGetInitialConversations = jest.fn(() => ({}));
   const mockGetComments = jest.fn(() => []);
   const mockHttp = httpServiceMock.createStartContract({ basePath: '/test' });
 
@@ -33,13 +39,13 @@ export const TestProvidersComponent: React.FC<Props> = ({ children }) => {
       <ThemeProvider theme={() => ({ eui: euiDarkVars, darkMode: true })}>
         <AssistantProvider
           actionTypeRegistry={actionTypeRegistry}
-          augmentMessageCodeBlocks={jest.fn()}
+          augmentMessageCodeBlocks={jest.fn().mockReturnValue([])}
           baseAllow={[]}
           baseAllowReplacement={[]}
           defaultAllow={[]}
           defaultAllowReplacement={[]}
           getComments={mockGetComments}
-          getInitialConversations={mockGetInitialConversations}
+          getInitialConversations={getInitialConversations}
           setConversations={jest.fn()}
           setDefaultAllow={jest.fn()}
           setDefaultAllowReplacement={jest.fn()}