[Security solution] Assistant telemetry conversation id fix (elastic#…

…173794)
kibanamachine · Dec 21, 2023 · 574ff80 · 574ff80
1 parent 80edac1
commit 574ff80
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 18 deletions.
diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_overlay/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_overlay/index.tsx
@@ -33,8 +33,7 @@ export const AssistantOverlay = React.memo(() => {
     WELCOME_CONVERSATION_TITLE
   );
   const [promptContextId, setPromptContextId] = useState<string | undefined>();
-  const { assistantTelemetry, setShowAssistantOverlay, localStorageLastConversationId } =
-    useAssistantContext();
+  const { assistantTelemetry, setShowAssistantOverlay, getConversationId } = useAssistantContext();
 
   // Bind `showAssistantOverlay` in SecurityAssistantContext to this modal instance
   const showOverlay = useCallback(
@@ -44,16 +43,18 @@ export const AssistantOverlay = React.memo(() => {
         promptContextId: pid,
         conversationId: cid,
       }: ShowAssistantOverlayProps) => {
+        const newConversationId = getConversationId(cid);
         if (so)
           assistantTelemetry?.reportAssistantInvoked({
-            conversationId: cid ?? 'unknown',
+            conversationId: newConversationId,
             invokedBy: 'click',
           });
+
         setIsModalVisible(so);
         setPromptContextId(pid);
-        setConversationId(cid);
+        setConversationId(newConversationId);
       },
-    [assistantTelemetry]
+    [assistantTelemetry, getConversationId]
   );
   useEffect(() => {
     setShowAssistantOverlay(showOverlay);
@@ -63,15 +64,15 @@ export const AssistantOverlay = React.memo(() => {
   const handleShortcutPress = useCallback(() => {
     // Try to restore the last conversation on shortcut pressed
     if (!isModalVisible) {
-      setConversationId(localStorageLastConversationId ?? WELCOME_CONVERSATION_TITLE);
+      setConversationId(getConversationId());
       assistantTelemetry?.reportAssistantInvoked({
         invokedBy: 'shortcut',
-        conversationId: localStorageLastConversationId ?? WELCOME_CONVERSATION_TITLE,
+        conversationId: getConversationId(),
       });
     }
 
     setIsModalVisible(!isModalVisible);
-  }, [assistantTelemetry, isModalVisible, localStorageLastConversationId]);
+  }, [assistantTelemetry, isModalVisible, getConversationId]);
 
   // Register keyboard listener to show the modal when cmd + ; is pressed
   const onKeyDown = useCallback(

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx
@@ -83,7 +83,7 @@ const AssistantComponent: React.FC<Props> = ({
     http,
     promptContexts,
     setLastConversationId,
-    localStorageLastConversationId,
+    getConversationId,
     title,
     allSystemPrompts,
   } = useAssistantContext();
@@ -113,12 +113,7 @@ const AssistantComponent: React.FC<Props> = ({
   );
 
   const [selectedConversationId, setSelectedConversationId] = useState<string>(
-    isAssistantEnabled
-      ? // if a conversationId has been provided, use that
-        // if not, check local storage
-        // last resort, go to welcome conversation
-        conversationId ?? localStorageLastConversationId ?? WELCOME_CONVERSATION_TITLE
-      : WELCOME_CONVERSATION_TITLE
+    isAssistantEnabled ? getConversationId(conversationId) : WELCOME_CONVERSATION_TITLE
   );
 
   useEffect(() => {

diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.test.tsx
@@ -12,7 +12,11 @@ import { AssistantProvider, useAssistantContext } from '.';
 import { httpServiceMock } from '@kbn/core-http-browser-mocks';
 import { actionTypeRegistryMock } from '@kbn/triggers-actions-ui-plugin/public/application/action_type_registry.mock';
 import { AssistantAvailability } from '../..';
+import { useLocalStorage } from 'react-use';
 
+jest.mock('react-use', () => ({
+  useLocalStorage: jest.fn().mockReturnValue(['456', jest.fn()]),
+}));
 const actionTypeRegistry = actionTypeRegistryMock.create();
 const mockGetInitialConversations = jest.fn(() => ({}));
 const mockGetComments = jest.fn(() => []);
@@ -70,4 +74,23 @@ describe('AssistantContext', () => {
 
     expect(mockHttp.fetch).toBeCalledWith(path);
   });
+
+  test('getConversationId defaults to provided id', async () => {
+    const { result } = renderHook(useAssistantContext, { wrapper: ContextWrapper });
+    const id = result.current.getConversationId('123');
+    expect(id).toEqual('123');
+  });
+
+  test('getConversationId uses local storage id when no id is provided ', async () => {
+    const { result } = renderHook(useAssistantContext, { wrapper: ContextWrapper });
+    const id = result.current.getConversationId();
+    expect(id).toEqual('456');
+  });
+
+  test('getConversationId defaults to Welcome when no local storage id and no id is provided ', async () => {
+    (useLocalStorage as jest.Mock).mockReturnValue([undefined, jest.fn()]);
+    const { result } = renderHook(useAssistantContext, { wrapper: ContextWrapper });
+    const id = result.current.getConversationId();
+    expect(id).toEqual('Welcome');
+  });
 });
diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx
@@ -13,6 +13,7 @@ import type { IToasts } from '@kbn/core-notifications-browser';
 import { ActionTypeRegistryContract } from '@kbn/triggers-actions-ui-plugin/public';
 import { useLocalStorage } from 'react-use';
 import type { DocLinksStart } from '@kbn/core-doc-links-browser';
+import { WELCOME_CONVERSATION_TITLE } from '../assistant/use_conversation/translations';
 import { updatePromptContexts } from './helpers';
 import type {
   PromptContext,
@@ -136,7 +137,7 @@ export interface UseAssistantContext {
   }) => EuiCommentProps[];
   http: HttpSetup;
   knowledgeBase: KnowledgeBaseConfig;
-  localStorageLastConversationId: string | undefined;
+  getConversationId: (id?: string) => string;
   promptContexts: Record<string, PromptContext>;
   modelEvaluatorEnabled: boolean;
   nameSpace: string;
@@ -292,6 +293,14 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
     [setConversations]
   );
 
+  const getConversationId = useCallback(
+    // if a conversationId has been provided, use that
+    // if not, check local storage
+    // last resort, go to welcome conversation
+    (id?: string) => id ?? localStorageLastConversationId ?? WELCOME_CONVERSATION_TITLE,
+    [localStorageLastConversationId]
+  );
+
   const value = useMemo(
     () => ({
       actionTypeRegistry,
@@ -334,7 +343,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       title,
       toasts,
       unRegisterPromptContext,
-      localStorageLastConversationId,
+      getConversationId,
       setLastConversationId: setLocalStorageLastConversationId,
     }),
     [
@@ -358,7 +367,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       getComments,
       http,
       localStorageKnowledgeBase,
-      localStorageLastConversationId,
+      getConversationId,
       localStorageQuickPrompts,
       localStorageSystemPrompts,
       modelEvaluatorEnabled,