[9.0] [SecuritySolution] Make last conversation local storage keys space aware (#214794)

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

@@ -14,6 +14,7 @@ import { css } from '@emotion/react';
 import { createGlobalStyle } from 'styled-components';
 import { ShowAssistantOverlayProps, useAssistantContext } from '../../assistant_context';
 import { Assistant, CONVERSATION_SIDE_PANEL_WIDTH } from '..';
+import { useAssistantLastConversation, useAssistantSpaceId } from '../use_space_aware_context';
 
 const isMac = navigator.platform.toLowerCase().indexOf('mac') >= 0;
 
@@ -29,12 +30,15 @@ export const UnifiedTimelineGlobalStyles = createGlobalStyle`
 `;
 
 export const AssistantOverlay = React.memo(() => {
+  const spaceId = useAssistantSpaceId();
   const [isModalVisible, setIsModalVisible] = useState(false);
   // Why is this named Title and not Id?
   const [conversationTitle, setConversationTitle] = useState<string | undefined>(undefined);
   const [promptContextId, setPromptContextId] = useState<string | undefined>();
-  const { assistantTelemetry, setShowAssistantOverlay, getLastConversationId } =
-    useAssistantContext();
+  const { assistantTelemetry, setShowAssistantOverlay } = useAssistantContext();
+  const { getLastConversationId } = useAssistantLastConversation({
+    spaceId,
+  });
 
   const [chatHistoryVisible, setChatHistoryVisible] = useState(false);
 
@@ -48,7 +52,6 @@ export const AssistantOverlay = React.memo(() => {
       }: ShowAssistantOverlayProps) => {
         const conversationId = getLastConversationId(cTitle);
         if (so) assistantTelemetry?.reportAssistantInvoked({ conversationId, invokedBy: 'click' });
-
         setIsModalVisible(so);
         setPromptContextId(pid);
         setConversationTitle(conversationId);

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

@@ -24,6 +24,7 @@ import { Conversation } from '../assistant_context/types';
 import * as all from './chat_send/use_chat_send';
 import { useConversation } from './use_conversation';
 import { AIConnector } from '../connectorland/connector_selector';
+import { localStorageLastConversationIdSubject$ } from './use_space_aware_context/use_last_conversation';
 
 jest.mock('../connectorland/use_load_connectors');
 jest.mock('../connectorland/connector_setup');
@@ -130,14 +131,16 @@ describe('Assistant', () => {
     } as unknown as DefinedUseQueryResult<Record<string, Conversation>, unknown>);
     jest
       .mocked(useLocalStorage)
-      .mockReturnValue([undefined, persistToLocalStorage] as unknown as ReturnType<
+      .mockReturnValue([mockData.welcome_id.id, persistToLocalStorage] as unknown as ReturnType<
         typeof useLocalStorage
       >);
     jest
       .mocked(useSessionStorage)
       .mockReturnValue([undefined, persistToSessionStorage] as unknown as ReturnType<
         typeof useSessionStorage
       >);
+
+    localStorageLastConversationIdSubject$.next(null);
   });
 
   describe('persistent storage', () => {

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

@@ -55,6 +55,7 @@ import {
   conversationContainsAnonymizedValues,
   conversationContainsContentReferences,
 } from './conversations/utils';
+import { useAssistantLastConversation, useAssistantSpaceId } from './use_space_aware_context';
 
 export const CONVERSATION_SIDE_PANEL_WIDTH = 220;
 
@@ -92,11 +93,9 @@ const AssistantComponent: React.FC<Props> = ({
     augmentMessageCodeBlocks,
     baseConversations,
     getComments,
-    getLastConversationId,
     http,
     promptContexts,
     currentUserAvatar,
-    setLastConversationId,
     contentReferencesVisible,
     showAnonymizedValues,
     setContentReferencesVisible,
@@ -133,6 +132,15 @@ const AssistantComponent: React.FC<Props> = ({
     http,
   });
   const defaultConnector = useMemo(() => getDefaultConnector(connectors), [connectors]);
+  const spaceId = useAssistantSpaceId();
+  const { setLastConversationId, getLastConversationId } = useAssistantLastConversation({
+    spaceId,
+  });
+  const lastConversationIdFromLocalStorage = useMemo(
+    () => getLastConversationId(),
+    [getLastConversationId]
+  );
+
   const {
     currentConversation,
     currentSystemPrompt,
@@ -147,14 +155,13 @@ const AssistantComponent: React.FC<Props> = ({
     conversations,
     defaultConnector,
     refetchCurrentUserConversations,
-    conversationId: getLastConversationId(conversationTitle),
+    conversationId: conversationTitle ?? lastConversationIdFromLocalStorage,
     mayUpdateConversations:
       isFetchedConnectors &&
       isFetchedCurrentUserConversations &&
       isFetchedPrompts &&
       Object.keys(conversations).length > 0,
   });
-
   const isInitialLoad = useMemo(() => {
     if (!isAssistantEnabled) {
       return false;

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

@@ -181,7 +181,6 @@ export const useAssistantOverlay = (
 
   useEffect(() => {
     unRegisterPromptContext(promptContextId); // a noop if the current prompt context id is not registered
-
     const newContext: PromptContext = {
       category: _category,
       description: _description,
@@ -191,7 +190,6 @@ export const useAssistantOverlay = (
       tooltip: _tooltip,
       replacements: _replacements ?? undefined,
     };
-
     registerPromptContext(newContext);
 
     return () => unRegisterPromptContext(promptContextId);

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

@@ -39,6 +39,7 @@ interface UseCurrentConversation {
     cId?: string;
     cTitle?: string;
     isStreamRefetch?: boolean;
+    silent?: boolean;
   }) => Promise<Conversation | undefined>;
   setCurrentConversation: Dispatch<SetStateAction<Conversation | undefined>>;
   setCurrentSystemPromptId: (promptId: string | undefined) => void;
@@ -109,13 +110,15 @@ export const useCurrentConversation = ({
    * @param cId - The conversation ID to refetch.
    * @param cTitle - The conversation title to refetch.
    * @param isStreamRefetch - Are we refetching because stream completed? If so retry several times to ensure the message has updated on the server
+   * @param silent - Should we show toasts on error
    */
   const refetchCurrentConversation = useCallback(
     async ({
       cId,
       cTitle,
       isStreamRefetch = false,
-    }: { cId?: string; cTitle?: string; isStreamRefetch?: boolean } = {}) => {
+      silent,
+    }: { cId?: string; cTitle?: string; isStreamRefetch?: boolean; silent?: boolean } = {}) => {
       if (cId === '' || (cTitle && !conversations[cTitle])) {
         return;
       }
@@ -124,7 +127,7 @@ export const useCurrentConversation = ({
         cId ?? (cTitle && conversations[cTitle].id) ?? currentConversation?.id;
 
       if (cConversationId) {
-        let updatedConversation = await getConversation(cConversationId);
+        let updatedConversation = await getConversation(cConversationId, silent);
         let retries = 0;
         const maxRetries = 5;
 
@@ -187,11 +190,7 @@ export const useCurrentConversation = ({
         setCurrentConversationId(cId);
       }
     },
-    [
-      initializeDefaultConversationWithConnector,
-      refetchCurrentUserConversations,
-      setCurrentConversationId,
-    ]
+    [initializeDefaultConversationWithConnector, refetchCurrentUserConversations]
   );
 
   // update currentConversation when conversations or currentConversationId update

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

@@ -0,0 +1,12 @@
+/*
+ * 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 { useAssistantSpaceId, AssistantSpaceIdProvider } from './use_space_id';
+import { useAssistantLastConversation, type LastConversation } from './use_last_conversation';
+
+export { useAssistantSpaceId, AssistantSpaceIdProvider, useAssistantLastConversation };
+export type { LastConversation };

x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/use_space_aware_context/use_last_conversation.test.tsx

@@ -0,0 +1,50 @@
+/*
+ * 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 useLocalStorage from 'react-use/lib/useLocalStorage';
+import { renderHook } from '@testing-library/react';
+import { useAssistantLastConversation } from './use_last_conversation';
+import {
+  DEFAULT_ASSISTANT_NAMESPACE,
+  LAST_CONVERSATION_ID_LOCAL_STORAGE_KEY,
+} from '../../assistant_context/constants';
+
+jest.mock('react-use/lib/useLocalStorage', () => jest.fn().mockReturnValue(['456', jest.fn()]));
+const spaceId = 'test';
+
+describe('useAssistantLastConversation', () => {
+  beforeEach(() => jest.clearAllMocks());
+
+  test('getLastConversationId defaults to provided id', async () => {
+    const { result } = renderHook(() => useAssistantLastConversation({ spaceId }));
+    const id = result.current.getLastConversationId('123');
+    expect(id).toEqual('123');
+  });
+
+  test('getLastConversationId uses local storage id when no id is provided ', async () => {
+    const { result } = renderHook(() => useAssistantLastConversation({ spaceId }));
+    const id = result.current.getLastConversationId();
+    expect(id).toEqual('456');
+  });
+
+  test('getLastConversationId 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(() => useAssistantLastConversation({ spaceId }));
+    const id = result.current.getLastConversationId();
+    expect(id).toEqual('Welcome');
+  });
+
+  describe.each([
+    {
+      expected: `${DEFAULT_ASSISTANT_NAMESPACE}.${LAST_CONVERSATION_ID_LOCAL_STORAGE_KEY}.${spaceId}`,
+    },
+  ])('useLocalStorage is called with keys with correct spaceId', ({ expected }) => {
+    test(`local storage key: ${expected}`, () => {
+      renderHook(() => useAssistantLastConversation({ spaceId }));
+      expect(useLocalStorage).toBeCalledWith(expected);
+    });
+  });
+});

x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/use_space_aware_context/use_last_conversation.tsx

@@ -0,0 +1,69 @@
+/*
+ * 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 { useCallback, useEffect, useMemo } from 'react';
+import useLocalStorage from 'react-use/lib/useLocalStorage';
+import { BehaviorSubject } from 'rxjs';
+import {
+  DEFAULT_ASSISTANT_NAMESPACE,
+  LAST_CONVERSATION_ID_LOCAL_STORAGE_KEY,
+} from '../../assistant_context/constants';
+import { WELCOME_CONVERSATION_TITLE } from '../use_conversation/translations';
+export interface LastConversation {
+  id: string;
+  title?: string;
+}
+
+export const localStorageLastConversationIdSubject$ = new BehaviorSubject<string | null>(null);
+
+export const useAssistantLastConversation = ({
+  nameSpace = DEFAULT_ASSISTANT_NAMESPACE,
+  spaceId,
+}: {
+  nameSpace?: string;
+  spaceId: string;
+}): {
+  getLastConversationId: (conversationId?: string) => string;
+  setLastConversationId: (conversationId?: string) => void;
+} => {
+  // Legacy fallback: used only if the new storage value is not yet set
+  const [localStorageLastConversationId, setLocalStorageLastConversationId] =
+    useLocalStorage<string>(`${nameSpace}.${LAST_CONVERSATION_ID_LOCAL_STORAGE_KEY}.${spaceId}`);
+
+  // Sync BehaviorSubject when localStorage changes
+  useEffect(() => {
+    if (localStorageLastConversationIdSubject$.getValue() !== localStorageLastConversationId) {
+      localStorageLastConversationIdSubject$.next(localStorageLastConversationId || null);
+    }
+  }, [localStorageLastConversationId]);
+
+  const getLastConversationId = useCallback(
+    // if a conversationId has been provided, use that
+    // if not, check local storage
+    // last resort, go to welcome conversation
+    (conversationId?: string) => {
+      return (
+        conversationId ??
+        localStorageLastConversationIdSubject$.getValue() ??
+        WELCOME_CONVERSATION_TITLE
+      );
+    },
+    []
+  );
+
+  const setLastConversationId = useCallback(
+    (conversationId?: string) => {
+      setLocalStorageLastConversationId(conversationId);
+      localStorageLastConversationIdSubject$.next(conversationId ?? null);
+    },
+    [setLocalStorageLastConversationId]
+  );
+
+  return useMemo(
+    () => ({ getLastConversationId, setLastConversationId }),
+    [getLastConversationId, setLastConversationId]
+  );
+};

x-pack/platform/packages/shared/kbn-elastic-assistant/impl/assistant/use_space_aware_context/use_space_id.tsx

@@ -0,0 +1,28 @@
+/*
+ * 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';
+
+interface UseSpaceIdContext {
+  spaceId: string;
+}
+interface SpaceIdProviderProps extends UseSpaceIdContext {
+  children: React.ReactNode;
+}
+
+const SpaceIdContext = React.createContext<UseSpaceIdContext | undefined>(undefined);
+
+export const AssistantSpaceIdProvider: React.FC<SpaceIdProviderProps> = ({ children, spaceId }) => {
+  return <SpaceIdContext.Provider value={{ spaceId }}>{children}</SpaceIdContext.Provider>;
+};
+
+export const useAssistantSpaceId = () => {
+  const context = React.useContext(SpaceIdContext);
+  if (context === undefined) {
+    throw new Error('useSpaceId must be used within a AssistantSpaceIdProvider');
+  }
+  return context.spaceId;
+};

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

@@ -8,12 +8,8 @@
 import { renderHook } from '@testing-library/react';
 
 import { useAssistantContext } from '.';
-import useLocalStorage from 'react-use/lib/useLocalStorage';
 import { TestProviders } from '../mock/test_providers/test_providers';
 
-jest.mock('react-use/lib/useLocalStorage', () => jest.fn().mockReturnValue(['456', jest.fn()]));
-jest.mock('react-use/lib/useSessionStorage', () => jest.fn().mockReturnValue(['456', jest.fn()]));
-
 describe('AssistantContext', () => {
   beforeEach(() => jest.clearAllMocks());
 
@@ -31,23 +27,4 @@ describe('AssistantContext', () => {
 
     expect(result.current.http.fetch).toBeCalledWith(path);
   });
-
-  test('getLastConversationId defaults to provided id', async () => {
-    const { result } = renderHook(useAssistantContext, { wrapper: TestProviders });
-    const id = result.current.getLastConversationId('123');
-    expect(id).toEqual('123');
-  });
-
-  test('getLastConversationId uses local storage id when no id is provided ', async () => {
-    const { result } = renderHook(useAssistantContext, { wrapper: TestProviders });
-    const id = result.current.getLastConversationId();
-    expect(id).toEqual('456');
-  });
-
-  test('getLastConversationId 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: TestProviders });
-    const id = result.current.getLastConversationId();
-    expect(id).toEqual('Welcome');
-  });
 });