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

## Summary

Issue and steps to reproduce:
#214114




https://github.com/user-attachments/assets/881fb082-a879-4816-b251-da3f2af77661



The best approach to fix this is to make local storage keys space aware.
In this use case, the current key is
`elasticAssistantDefault.lastSelectedConversation`.

Ideally it should be e.g.:
`elasticAssistantDefault.lastSelectedConversation.{spaceId}`

To retrieve spaceId properly, we have to make sure spaceId has been
available when reading the local storage. In other words, the spaceId
cannot be null, undefined, or a fallback value when accessing it.

To achieve that, we have to only render the assistant overlay after
spaceId is available and remove existing spaceId from the
AssistantContext. The reason I removed spaceId from AssistantContext is
because it is at the top of the rendering tree and can block the entire
app if waiting for the spaceId becomes available there.

`useAssistantSpaceId` comes from `SpaceIdContext` which should be render
only when spaceId is available, should be safe to access the exact
spaceId inside Assistant overly.


### Checklist

Check the PR satisfies following conditions. 

Reviewers should verify this PR satisfies this list as well.

- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

---------

Co-authored-by: kibanamachine <[email protected]>

Author: angorayc

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

@@ -12,12 +12,13 @@ import useEvent from 'react-use/lib/useEvent';
 import { css } from '@emotion/react';
 
 import { createGlobalStyle } from 'styled-components';
-import {
-  LastConversation,
-  ShowAssistantOverlayProps,
-  useAssistantContext,
-} from '../../assistant_context';
+import { ShowAssistantOverlayProps, useAssistantContext } from '../../assistant_context';
 import { Assistant, CONVERSATION_SIDE_PANEL_WIDTH } from '..';
+import {
+  useAssistantLastConversation,
+  useAssistantSpaceId,
+  type LastConversation,
+} from '../use_space_aware_context';
 
 const isMac = navigator.platform.toLowerCase().indexOf('mac') >= 0;
 
@@ -33,14 +34,15 @@ export const UnifiedTimelineGlobalStyles = createGlobalStyle`
 `;
 
 export const AssistantOverlay = React.memo(() => {
+  const spaceId = useAssistantSpaceId();
   const [isModalVisible, setIsModalVisible] = useState(false);
   // id if the conversation exists in the data stream, title if it's a new conversation
   const [lastConversation, setSelectedConversation] = useState<LastConversation | undefined>(
     undefined
   );
   const [promptContextId, setPromptContextId] = useState<string | undefined>();
-  const { assistantTelemetry, setShowAssistantOverlay, getLastConversation } =
-    useAssistantContext();
+  const { assistantTelemetry, setShowAssistantOverlay } = useAssistantContext();
+  const { getLastConversation } = useAssistantLastConversation({ spaceId });
 
   const [chatHistoryVisible, setChatHistoryVisible] = useState(false);
 

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

@@ -18,6 +18,7 @@ import { useConversation } from '../use_conversation';
 import { getCombinedMessage } from '../prompt/helpers';
 import { Conversation, useAssistantContext } from '../../..';
 import { getMessageFromRawResponse } from '../helpers';
+import { useAssistantSpaceId, useAssistantLastConversation } from '../use_space_aware_context';
 
 export interface UseChatSendProps {
   currentConversation?: Conversation;
@@ -55,8 +56,9 @@ export const useChatSend = ({
     assistantTelemetry,
     toasts,
     assistantAvailability: { isAssistantEnabled },
-    setLastConversation,
   } = useAssistantContext();
+  const spaceId = useAssistantSpaceId();
+  const { setLastConversation } = useAssistantLastConversation({ spaceId });
   const [userPrompt, setUserPrompt] = useState<string | null>(null);
 
   const { isLoading, sendMessage, abortStream } = useSendMessage();

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

@@ -149,7 +149,7 @@ describe('Assistant', () => {
       .mockReturnValue(defaultFetchUserConversations as unknown as FetchCurrentUserConversations);
     jest
       .mocked(useLocalStorage)
-      .mockReturnValue([undefined, persistToLocalStorage] as unknown as ReturnType<
+      .mockReturnValue([mockData.welcome_id, persistToLocalStorage] as unknown as ReturnType<
         typeof useLocalStorage
       >);
     jest

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

@@ -37,7 +37,7 @@ import { useChatSend } from './chat_send/use_chat_send';
 import { ChatSend } from './chat_send';
 import { getDefaultConnector } from './helpers';
 
-import { LastConversation, useAssistantContext } from '../assistant_context';
+import { useAssistantContext } from '../assistant_context';
 import { ContextPills } from './context_pills';
 import { getNewSelectedPromptContext } from '../data_anonymization/get_new_selected_prompt_context';
 import type { PromptContext, SelectedPromptContext } from './prompt_context/types';
@@ -53,6 +53,11 @@ import {
   conversationContainsAnonymizedValues,
   conversationContainsContentReferences,
 } from './conversations/utils';
+import {
+  LastConversation,
+  useAssistantLastConversation,
+  useAssistantSpaceId,
+} from './use_space_aware_context';
 
 export const CONVERSATION_SIDE_PANEL_WIDTH = 220;
 
@@ -90,12 +95,9 @@ const AssistantComponent: React.FC<Props> = ({
     currentAppId,
     augmentMessageCodeBlocks,
     getComments,
-    getLastConversation,
     http,
     promptContexts,
     currentUserAvatar,
-    setLastConversation,
-    spaceId,
     contentReferencesVisible,
     showAnonymizedValues,
     setContentReferencesVisible,
@@ -134,6 +136,13 @@ const AssistantComponent: React.FC<Props> = ({
     http,
   });
   const defaultConnector = useMemo(() => getDefaultConnector(connectors), [connectors]);
+  const spaceId = useAssistantSpaceId();
+  const { getLastConversation, setLastConversation } = useAssistantLastConversation({ spaceId });
+  const lastConversationFromLocalStorage = useMemo(
+    () => getLastConversation(),
+    [getLastConversation]
+  );
+
   const {
     currentConversation,
     currentSystemPrompt,
@@ -151,7 +160,7 @@ const AssistantComponent: React.FC<Props> = ({
     defaultConnector,
     spaceId,
     refetchCurrentUserConversations,
-    lastConversation: lastConversation ?? getLastConversation(lastConversation),
+    lastConversation: lastConversation ?? lastConversationFromLocalStorage,
     mayUpdateConversations:
       isFetchedConnectors && isFetchedCurrentUserConversations && isFetchedPrompts,
     setLastConversation,

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

@@ -229,12 +229,15 @@ describe('useCurrentConversation', () => {
       }),
     });
 
+    expect(mockUseConversation.getConversation).toHaveBeenCalledWith(mockData.welcome_id.id, true);
+
     await act(async () => {
       await result.current.handleOnConversationSelected({
         cId: conversationId,
       });
     });
 
+    expect(mockUseConversation.getConversation).toHaveBeenCalledWith(conversationId, undefined);
     expect(result.current.currentConversation).toEqual(conversation);
     expect(result.current.currentSystemPrompt?.id).toBe('something-crazy');
   });

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

@@ -14,13 +14,13 @@ import {
 } from '@tanstack/react-query';
 import { ApiConfig, PromptResponse } from '@kbn/elastic-assistant-common';
 import useLocalStorage from 'react-use/lib/useLocalStorage';
-import { LastConversation } from '../../assistant_context';
 import { FetchConversationsResponse } from '../api';
 import { AIConnector } from '../../connectorland/connector_selector';
 import { getDefaultNewSystemPrompt, getDefaultSystemPrompt } from '../use_conversation/helpers';
 import { useConversation } from '../use_conversation';
 import { sleep } from '../helpers';
 import { Conversation } from '../../..';
+import type { LastConversation } from '../use_space_aware_context';
 
 export interface Props {
   allSystemPrompts: PromptResponse[];
@@ -34,7 +34,7 @@ export interface Props {
   refetchCurrentUserConversations: <TPageData>(
     options?: RefetchOptions & RefetchQueryFilters<TPageData>
   ) => Promise<QueryObserverResult<InfiniteData<FetchConversationsResponse>, unknown>>;
-  setLastConversation: Dispatch<SetStateAction<LastConversation | undefined>>;
+  setLastConversation: (lastConversation: LastConversation) => void;
 }
 
 interface UseCurrentConversation {
@@ -125,14 +125,18 @@ export const useCurrentConversation = ({
    * @param isStreamRefetch - Are we refetching because stream completed? If so retry several times to ensure the message has updated on the server
    */
   const refetchCurrentConversation = useCallback(
-    async ({ cId, isStreamRefetch = false }: { cId?: string; isStreamRefetch?: boolean } = {}) => {
+    async ({
+      cId,
+      isStreamRefetch = false,
+      silent,
+    }: { cId?: string; isStreamRefetch?: boolean; silent?: boolean } = {}) => {
       if (cId === '') {
         return;
       }
       const cConversationId = cId ?? currentConversation?.id;
 
       if (cConversationId) {
-        let updatedConversation = await getConversation(cConversationId);
+        let updatedConversation = await getConversation(cConversationId, silent);
         let retries = 0;
         const maxRetries = 5;
 
@@ -203,7 +207,17 @@ export const useCurrentConversation = ({
   );
 
   const handleOnConversationSelected = useCallback(
-    async ({ cId, cTitle, apiConfig }: { apiConfig?: ApiConfig; cId: string; cTitle?: string }) => {
+    async ({
+      cId,
+      cTitle,
+      apiConfig,
+      silent,
+    }: {
+      apiConfig?: ApiConfig;
+      cId: string;
+      cTitle?: string;
+      silent?: boolean;
+    }) => {
       if (cId === '') {
         if (
           currentAppId === 'securitySolutionUI' &&
@@ -228,7 +242,7 @@ export const useCurrentConversation = ({
       }
       // refetch will set the currentConversation
       try {
-        await refetchCurrentConversation({ cId });
+        await refetchCurrentConversation({ cId, silent });
         setLastConversation({
           id: cId,
         });
@@ -249,7 +263,11 @@ export const useCurrentConversation = ({
 
   useEffect(() => {
     if (!mayUpdateConversations || !!currentConversation) return;
-    handleOnConversationSelected({ cId: lastConversation.id, cTitle: lastConversation.title });
+    handleOnConversationSelected({
+      cId: lastConversation.id,
+      cTitle: lastConversation.title,
+      silent: true,
+    });
   }, [lastConversation, handleOnConversationSelected, currentConversation, mayUpdateConversations]);
 
   const handleOnConversationDeleted = useCallback(

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,63 @@
+/*
+ * 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,
+  LAST_SELECTED_CONVERSATION_LOCAL_STORAGE_KEY,
+} from '../../assistant_context/constants';
+
+jest.mock('react-use/lib/useLocalStorage', () =>
+  jest.fn().mockReturnValue([{ id: '456' }, jest.fn()])
+);
+const spaceId = 'test';
+
+describe('useAssistantLastConversation', () => {
+  beforeEach(() => jest.clearAllMocks());
+
+  test('getLastConversation defaults to provided id', () => {
+    const { result } = renderHook(() => useAssistantLastConversation({ spaceId }));
+    const id = result.current.getLastConversation({ id: '123' });
+    expect(id).toEqual({ id: '123' });
+  });
+
+  test('getLastConversation uses local storage id when no id is provided ', () => {
+    const { result } = renderHook(() => useAssistantLastConversation({ spaceId }));
+    const id = result.current.getLastConversation();
+    expect(id).toEqual({ id: '456' });
+  });
+
+  test('getLastConversation defaults to empty id when no local storage id and no id is provided ', () => {
+    (useLocalStorage as jest.Mock).mockReturnValue([undefined, jest.fn()]);
+    const { result } = renderHook(() => useAssistantLastConversation({ spaceId }));
+    const id = result.current.getLastConversation();
+    expect(id).toEqual({ id: '' });
+  });
+
+  test('getLastConversation defaults to empty id when title is provided and preserves title', () => {
+    (useLocalStorage as jest.Mock).mockReturnValue([undefined, jest.fn()]);
+    const { result } = renderHook(() => useAssistantLastConversation({ spaceId }));
+    const id = result.current.getLastConversation({ title: 'something' });
+    expect(id).toEqual({ id: '', title: 'something' });
+  });
+
+  describe.each([
+    {
+      expected: `${DEFAULT_ASSISTANT_NAMESPACE}.${LAST_CONVERSATION_ID_LOCAL_STORAGE_KEY}.${spaceId}`,
+    },
+    {
+      expected: `${DEFAULT_ASSISTANT_NAMESPACE}.${LAST_SELECTED_CONVERSATION_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);
+    });
+  });
+});