[Security Solution] Store last conversation in localstorage #6993 (#1…

…61373)

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

@@ -36,7 +36,7 @@ export const AssistantOverlay = React.memo<Props>(({ isAssistantEnabled }) => {
     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(
@@ -56,15 +56,25 @@ export const AssistantOverlay = React.memo<Props>(({ isAssistantEnabled }) => {
     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,163 @@
+/*
+ * 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 type { 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 isAssistantEnabled />
+    </TestProviders>
+  );
+
+describe('Assistant', () => {
+  beforeAll(() => {
+    jest.mocked(useConnectorSetup).mockReturnValue({
+      comments: [],
+      prompt: <></>,
+    });
+
+    jest.mocked(PromptEditor).mockReturnValue(null);
+    jest.mocked(QuickPrompts).mockReturnValue(null);
+  });
+
+  let persistToLocalStorage: jest.Mock;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+    persistToLocalStorage = jest.fn();
+
+    jest
+      .mocked(useLocalStorage)
+      .mockReturnValue([undefined, 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');
+    });
+
+    it('should not persist the conversation id to local storage when excludeFromLastConversationStorage flag is indicated', async () => {
+      const connectors: unknown[] = [{}];
+
+      jest.mocked(useLoadConnectors).mockReturnValue({
+        isSuccess: true,
+        data: connectors,
+      } as unknown as UseQueryResult<ActionConnector[], IHttpFetchError>);
+
+      const { getByLabelText } = render(
+        <TestProviders
+          getInitialConversations={() => ({
+            [WELCOME_CONVERSATION_TITLE]: {
+              id: WELCOME_CONVERSATION_TITLE,
+              messages: [],
+              apiConfig: {},
+            },
+            [MOCK_CONVERSATION_TITLE]: {
+              id: MOCK_CONVERSATION_TITLE,
+              messages: [],
+              apiConfig: {},
+              excludeFromLastConversationStorage: true,
+            },
+          })}
+        >
+          <Assistant isAssistantEnabled />
+        </TestProviders>
+      );
+
+      expect(persistToLocalStorage).toHaveBeenCalled();
+
+      expect(persistToLocalStorage).toHaveBeenLastCalledWith(WELCOME_CONVERSATION_TITLE);
+
+      const previousConversationButton = getByLabelText('Previous conversation');
+
+      expect(previousConversationButton).toBeInTheDocument();
+
+      await act(async () => {
+        fireEvent.click(previousConversationButton);
+      });
+      expect(persistToLocalStorage).toHaveBeenLastCalledWith(WELCOME_CONVERSATION_TITLE);
+    });
+  });
+
+  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

@@ -81,6 +81,7 @@ const AssistantComponent: React.FC<Props> = ({
     getComments,
     http,
     promptContexts,
+    setLastConversationId,
     title,
     allSystemPrompts,
   } = useAssistantContext();
@@ -136,7 +137,11 @@ const AssistantComponent: React.FC<Props> = ({
         };
   }, [conversations, isAssistantEnabled, 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(
     () =>
@@ -145,6 +150,18 @@ 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('');
+    }
+
+    if (!currentConversation.excludeFromLastConversationStorage) {
+      setLastConversationId(currentConversation.id);
+    }
+  }, [areConnectorsFetched, connectors?.length, currentConversation, setLastConversationId]);
+
   const isWelcomeSetup = (connectors?.length ?? 0) === 0;
   const isDisabled = isWelcomeSetup || !isAssistantEnabled;
 

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

@@ -28,6 +28,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';
@@ -99,6 +100,7 @@ interface UseAssistantContext {
     showAnonymizedValues: boolean;
   }) => EuiCommentProps[];
   http: HttpSetup;
+  localStorageLastConversationId: string | undefined;
   promptContexts: Record<string, PromptContext>;
   nameSpace: string;
   registerPromptContext: RegisterPromptContext;
@@ -107,6 +109,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;
@@ -158,6 +161,9 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
     setLocalStorageSystemPrompts(baseSystemPrompts);
   }, [baseSystemPrompts, setLocalStorageSystemPrompts]);
 
+  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.
    */
@@ -259,6 +265,8 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       showAssistantOverlay,
       title,
       unRegisterPromptContext,
+      localStorageLastConversationId,
+      setLastConversationId: setLocalStorageLastConversationId,
     }),
     [
       actionTypeRegistry,
@@ -275,6 +283,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       docLinks,
       getComments,
       http,
+      localStorageLastConversationId,
       localStorageQuickPrompts,
       localStorageSystemPrompts,
       nameSpace,
@@ -283,6 +292,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       registerPromptContext,
       setDefaultAllow,
       setDefaultAllowReplacement,
+      setLocalStorageLastConversationId,
       setLocalStorageQuickPrompts,
       setLocalStorageSystemPrompts,
       showAssistantOverlay,

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

@@ -53,9 +53,5 @@ export interface Conversation {
   replacements?: Record<string, string>;
   theme?: ConversationTheme;
   isDefault?: boolean;
-}
-
-export interface OpenAIConfig {
-  temperature: number;
-  model: string;
+  excludeFromLastConversationStorage?: boolean;
 }

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

@@ -14,17 +14,24 @@ 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();
+window.HTMLElement.prototype.scrollIntoView = 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,7 +40,7 @@ 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={[]}
@@ -43,7 +50,7 @@ export const TestProvidersComponent: React.FC<Props> = ({ children }) => {
             DOC_LINK_VERSION: 'current',
           }}
           getComments={mockGetComments}
-          getInitialConversations={mockGetInitialConversations}
+          getInitialConversations={getInitialConversations}
           setConversations={jest.fn()}
           setDefaultAllow={jest.fn()}
           setDefaultAllowReplacement={jest.fn()}

x-pack/packages/kbn-elastic-assistant/tsconfig.json

@@ -28,5 +28,6 @@
     "@kbn/i18n-react",
     "@kbn/ui-theme",
     "@kbn/core-doc-links-browser",
+    "@kbn/core",
   ]
 }

x-pack/plugins/security_solution/public/assistant/content/conversations/index.tsx

@@ -45,6 +45,7 @@ export const BASE_SECURITY_CONVERSATIONS: Record<string, Conversation> = {
     apiConfig: {},
   },
   [TIMELINE_CONVERSATION_TITLE]: {
+    excludeFromLastConversationStorage: true,
     id: TIMELINE_CONVERSATION_TITLE,
     isDefault: true,
     messages: [],