[Security Solution] AI Assistant telemetry

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

@@ -0,0 +1,85 @@
+/*
+ * 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 { fireEvent, render } from '@testing-library/react';
+import { AssistantOverlay } from '.';
+import { TestProviders } from '../../mock/test_providers/test_providers';
+
+const reportAssistantInvoked = jest.fn();
+const assistantTelemetry = {
+  reportAssistantInvoked,
+  reportAssistantMessageSent: () => {},
+  reportAssistantQuickPrompt: () => {},
+};
+describe('AssistantOverlay', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+  it('renders when isAssistantEnabled prop is true and keyboard shortcut is pressed', () => {
+    const { getByTestId } = render(
+      <TestProviders providerContext={{ assistantTelemetry }}>
+        <AssistantOverlay isAssistantEnabled={true} />
+      </TestProviders>
+    );
+    fireEvent.keyDown(document, { key: ';', ctrlKey: true });
+    const modal = getByTestId('ai-assistant-modal');
+    expect(modal).toBeInTheDocument();
+  });
+
+  it('modal closes when close button is clicked', () => {
+    const { getByLabelText, queryByTestId } = render(
+      <TestProviders>
+        <AssistantOverlay isAssistantEnabled={true} />
+      </TestProviders>
+    );
+    fireEvent.keyDown(document, { key: ';', ctrlKey: true });
+    const closeButton = getByLabelText('Closes this modal window');
+    fireEvent.click(closeButton);
+    const modal = queryByTestId('ai-assistant-modal');
+    expect(modal).not.toBeInTheDocument();
+  });
+
+  it('Assistant invoked from shortcut tracking happens on modal open only (not close)', () => {
+    render(
+      <TestProviders providerContext={{ assistantTelemetry }}>
+        <AssistantOverlay isAssistantEnabled={true} />
+      </TestProviders>
+    );
+    fireEvent.keyDown(document, { key: ';', ctrlKey: true });
+    expect(reportAssistantInvoked).toHaveBeenCalledTimes(1);
+    expect(reportAssistantInvoked).toHaveBeenCalledWith({
+      invokedBy: 'shortcut',
+      conversationId: 'Welcome',
+    });
+    fireEvent.keyDown(document, { key: ';', ctrlKey: true });
+    expect(reportAssistantInvoked).toHaveBeenCalledTimes(1);
+  });
+
+  it('modal closes when shortcut is pressed and modal is already open', () => {
+    const { queryByTestId } = render(
+      <TestProviders>
+        <AssistantOverlay isAssistantEnabled={true} />
+      </TestProviders>
+    );
+    fireEvent.keyDown(document, { key: ';', ctrlKey: true });
+    fireEvent.keyDown(document, { key: ';', ctrlKey: true });
+    const modal = queryByTestId('ai-assistant-modal');
+    expect(modal).not.toBeInTheDocument();
+  });
+
+  it('modal does not open when incorrect shortcut is pressed', () => {
+    const { queryByTestId } = render(
+      <TestProviders>
+        <AssistantOverlay isAssistantEnabled={true} />
+      </TestProviders>
+    );
+    fireEvent.keyDown(document, { key: 'a', ctrlKey: true });
+    const modal = queryByTestId('ai-assistant-modal');
+    expect(modal).not.toBeInTheDocument();
+  });
+});

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

@@ -36,7 +36,8 @@ export const AssistantOverlay = React.memo<Props>(({ isAssistantEnabled }) => {
     WELCOME_CONVERSATION_TITLE
   );
   const [promptContextId, setPromptContextId] = useState<string | undefined>();
-  const { setShowAssistantOverlay, localStorageLastConversationId } = useAssistantContext();
+  const { assistantTelemetry, setShowAssistantOverlay, localStorageLastConversationId } =
+    useAssistantContext();
 
   // Bind `showAssistantOverlay` in SecurityAssistantContext to this modal instance
   const showOverlay = useCallback(
@@ -46,11 +47,16 @@ export const AssistantOverlay = React.memo<Props>(({ isAssistantEnabled }) => {
         promptContextId: pid,
         conversationId: cid,
       }: ShowAssistantOverlayProps) => {
+        if (so)
+          assistantTelemetry?.reportAssistantInvoked({
+            conversationId: cid ?? 'unknown',
+            invokedBy: 'click',
+          });
         setIsModalVisible(so);
         setPromptContextId(pid);
         setConversationId(cid);
       },
-    [setIsModalVisible]
+    [assistantTelemetry]
   );
   useEffect(() => {
     setShowAssistantOverlay(showOverlay);
@@ -61,10 +67,14 @@ export const AssistantOverlay = React.memo<Props>(({ isAssistantEnabled }) => {
     // Try to restore the last conversation on shortcut pressed
     if (!isModalVisible) {
       setConversationId(localStorageLastConversationId ?? WELCOME_CONVERSATION_TITLE);
+      assistantTelemetry?.reportAssistantInvoked({
+        invokedBy: 'shortcut',
+        conversationId: localStorageLastConversationId ?? WELCOME_CONVERSATION_TITLE,
+      });
     }
 
     setIsModalVisible(!isModalVisible);
-  }, [isModalVisible, localStorageLastConversationId]);
+  }, [assistantTelemetry, isModalVisible, localStorageLastConversationId]);
 
   // Register keyboard listener to show the modal when cmd + ; is pressed
   const onKeyDown = useCallback(

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

@@ -46,10 +46,10 @@ const getInitialConversations = (): Record<string, Conversation> => ({
   },
 });
 
-const renderAssistant = () =>
+const renderAssistant = (extraProps = {}) =>
   render(
     <TestProviders getInitialConversations={getInitialConversations}>
-      <Assistant isAssistantEnabled />
+      <Assistant isAssistantEnabled {...extraProps} />
     </TestProviders>
   );
 
@@ -143,6 +143,22 @@ describe('Assistant', () => {
       });
       expect(persistToLocalStorage).toHaveBeenLastCalledWith(WELCOME_CONVERSATION_TITLE);
     });
+    it('should call the setConversationId callback if it is defined and the conversation id changes', async () => {
+      const connectors: unknown[] = [{}];
+      const setConversationId = jest.fn();
+      jest.mocked(useLoadConnectors).mockReturnValue({
+        isSuccess: true,
+        data: connectors,
+      } as unknown as UseQueryResult<ActionConnector[], IHttpFetchError>);
+
+      renderAssistant({ setConversationId });
+
+      await act(async () => {
+        fireEvent.click(screen.getByLabelText('Previous conversation'));
+      });
+
+      expect(setConversationId).toHaveBeenLastCalledWith('electric sheep');
+    });
   });
 
   describe('when no connectors are loaded', () => {

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

@@ -5,7 +5,16 @@
  * 2.0.
  */
 
-import React, { useCallback, useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react';
+import React, {
+  Dispatch,
+  SetStateAction,
+  useCallback,
+  useEffect,
+  useLayoutEffect,
+  useMemo,
+  useRef,
+  useState,
+} from 'react';
 import {
   EuiFlexGroup,
   EuiFlexItem,
@@ -46,6 +55,7 @@ export interface Props {
   promptContextId?: string;
   shouldRefocusPrompt?: boolean;
   showTitle?: boolean;
+  setConversationId?: Dispatch<SetStateAction<string>>;
 }
 
 /**
@@ -58,9 +68,11 @@ const AssistantComponent: React.FC<Props> = ({
   promptContextId = '',
   shouldRefocusPrompt = false,
   showTitle = true,
+  setConversationId,
 }) => {
   const {
     actionTypeRegistry,
+    assistantTelemetry,
     augmentMessageCodeBlocks,
     conversations,
     defaultAllow,
@@ -112,6 +124,12 @@ const AssistantComponent: React.FC<Props> = ({
       : WELCOME_CONVERSATION_TITLE
   );
 
+  useEffect(() => {
+    if (setConversationId) {
+      setConversationId(selectedConversationId);
+    }
+  }, [selectedConversationId, setConversationId]);
+
   const currentConversation = useMemo(
     () =>
       conversations[selectedConversationId] ??
@@ -396,6 +414,16 @@ const AssistantComponent: React.FC<Props> = ({
     return chatbotComments;
   }, [connectorComments, isDisabled, chatbotComments]);
 
+  const trackPrompt = useCallback(
+    (promptTitle: string) => {
+      assistantTelemetry?.reportAssistantQuickPrompt({
+        conversationId: selectedConversationId,
+        promptTitle,
+      });
+    },
+    [assistantTelemetry, selectedConversationId]
+  );
+
   return (
     <>
       <EuiModalHeader
@@ -485,6 +513,7 @@ const AssistantComponent: React.FC<Props> = ({
           <QuickPrompts
             setInput={setUserPrompt}
             setIsSettingsModalVisible={setIsSettingsModalVisible}
+            trackPrompt={trackPrompt}
           />
         )}
       </EuiModalFooter>

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

@@ -0,0 +1,75 @@
+/*
+ * 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 { fireEvent, render } from '@testing-library/react';
+import { QuickPrompts } from './quick_prompts';
+import { TestProviders } from '../../mock/test_providers/test_providers';
+import { MOCK_QUICK_PROMPTS } from '../../mock/quick_prompt';
+import { QUICK_PROMPTS_TAB } from '../settings/assistant_settings';
+
+const setInput = jest.fn();
+const setIsSettingsModalVisible = jest.fn();
+const trackPrompt = jest.fn();
+const testProps = {
+  setInput,
+  setIsSettingsModalVisible,
+  trackPrompt,
+};
+const setSelectedSettingsTab = jest.fn();
+const mockUseAssistantContext = {
+  setSelectedSettingsTab,
+  promptContexts: {},
+  allQuickPrompts: MOCK_QUICK_PROMPTS,
+};
+
+const testTitle = 'SPL_QUERY_CONVERSION_TITLE';
+const testPrompt = 'SPL_QUERY_CONVERSION_PROMPT';
+const customTitle = 'A_CUSTOM_OPTION';
+
+jest.mock('../../assistant_context', () => ({
+  ...jest.requireActual('../../assistant_context'),
+  useAssistantContext: () => mockUseAssistantContext,
+}));
+
+describe('QuickPrompts', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+  it('onClickAddQuickPrompt calls setInput with the prompt, and trackPrompt with the prompt title', () => {
+    const { getByText } = render(
+      <TestProviders>
+        <QuickPrompts {...testProps} />
+      </TestProviders>
+    );
+    fireEvent.click(getByText(testTitle));
+
+    expect(setInput).toHaveBeenCalledWith(testPrompt);
+    expect(trackPrompt).toHaveBeenCalledWith(testTitle);
+  });
+  it('onClickAddQuickPrompt calls trackPrompt with "Custom" when isDefault=false prompt is chosen', () => {
+    const { getByText } = render(
+      <TestProviders>
+        <QuickPrompts {...testProps} />
+      </TestProviders>
+    );
+    fireEvent.click(getByText(customTitle));
+
+    expect(trackPrompt).toHaveBeenCalledWith('Custom');
+  });
+
+  it('clicking "Add quick prompt" button opens the settings modal', () => {
+    const { getByTestId } = render(
+      <TestProviders>
+        <QuickPrompts {...testProps} />
+      </TestProviders>
+    );
+    fireEvent.click(getByTestId('addQuickPrompt'));
+    expect(setIsSettingsModalVisible).toHaveBeenCalledWith(true);
+    expect(setSelectedSettingsTab).toHaveBeenCalledWith(QUICK_PROMPTS_TAB);
+  });
+});