[Obs AI Assistant] Persist flyout mode preferences (#205366)

Persist flyout mode settings (docked and/or open/closed) across page
refreshes. I.e., if a user selects docked mode while the flyout is open,
and refreshes the page, the flyout will open on page load, in docked
mode.

x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_flyout.tsx

@@ -44,6 +44,7 @@ export function ChatFlyout({
   initialTitle,
   initialMessages,
   initialFlyoutPositionMode,
+  onFlyoutPositionModeChange,
   isOpen,
   onClose,
   navigateToConversation,
@@ -52,6 +53,7 @@ export function ChatFlyout({
   initialTitle: string;
   initialMessages: Message[];
   initialFlyoutPositionMode?: FlyoutPositionMode;
+  onFlyoutPositionModeChange?: (next: FlyoutPositionMode) => void;
   isOpen: boolean;
   onClose: () => void;
   navigateToConversation?: (conversationId?: string) => void;
@@ -137,6 +139,7 @@ export function ChatFlyout({
 
   const handleToggleFlyoutPositionMode = (newFlyoutPositionMode: FlyoutPositionMode) => {
     setFlyoutPositionMode(newFlyoutPositionMode);
+    onFlyoutPositionModeChange?.(newFlyoutPositionMode);
   };
 
   return isOpen ? (

x-pack/solutions/observability/plugins/observability_ai_assistant_app/public/components/nav_control/index.tsx

@@ -12,7 +12,12 @@ import { v4 } from 'uuid';
 import useObservable from 'react-use/lib/useObservable';
 import { i18n } from '@kbn/i18n';
 import { CoreStart } from '@kbn/core-lifecycle-browser';
-import { AIAssistantAppService, useAIAssistantAppService, ChatFlyout } from '@kbn/ai-assistant';
+import {
+  AIAssistantAppService,
+  useAIAssistantAppService,
+  ChatFlyout,
+  FlyoutPositionMode,
+} from '@kbn/ai-assistant';
 import { AssistantIcon } from '@kbn/ai-assistant-icon';
 import { useKibana } from '../../hooks/use_kibana';
 import { useTheme } from '../../hooks/use_theme';
@@ -20,6 +25,7 @@ import { useNavControlScreenContext } from '../../hooks/use_nav_control_screen_c
 import { SharedProviders } from '../../utils/shared_providers';
 import { ObservabilityAIAssistantAppPluginStartDependencies } from '../../types';
 import { useNavControlScope } from '../../hooks/use_nav_control_scope';
+import { useLocalStorage } from '../../hooks/use_local_storage';
 
 interface NavControlWithProviderDeps {
   appService: AIAssistantAppService;
@@ -62,11 +68,21 @@ export function NavControl({ isServerless }: { isServerless?: boolean }) {
     },
   } = useKibana();
 
-  const [hasBeenOpened, setHasBeenOpened] = useState(false);
-
   useNavControlScreenContext();
   useNavControlScope();
 
+  const [flyoutSettings, setFlyoutSettings] = useLocalStorage(
+    'observabilityAIAssistant.flyoutSettings',
+    {
+      mode: FlyoutPositionMode.OVERLAY,
+      isOpen: false,
+    }
+  );
+
+  const [isOpen, setIsOpen] = useState(flyoutSettings.isOpen);
+  const [hasBeenOpened, setHasBeenOpened] = useState(isOpen);
+  const keyRef = useRef(v4());
+
   const chatService = useAbortableAsync(
     ({ signal }) => {
       return hasBeenOpened
@@ -90,21 +106,18 @@ export function NavControl({ isServerless }: { isServerless?: boolean }) {
     [service, hasBeenOpened, notifications.toasts]
   );
 
-  const [isOpen, setIsOpen] = useState(false);
-
-  const keyRef = useRef(v4());
-
   useEffect(() => {
     const conversationSubscription = service.conversations.predefinedConversation$.subscribe(() => {
       keyRef.current = v4();
       setHasBeenOpened(true);
+      setFlyoutSettings((prev) => ({ ...prev, isOpen: true }));
       setIsOpen(true);
     });
 
     return () => {
       conversationSubscription.unsubscribe();
     };
-  }, [service.conversations.predefinedConversation$]);
+  }, [service.conversations.predefinedConversation$, setFlyoutSettings]);
 
   const { messages, title, hideConversationList } = useObservable(
     service.conversations.predefinedConversation$
@@ -186,9 +199,14 @@ export function NavControl({ isServerless }: { isServerless?: boolean }) {
             isOpen={isOpen}
             initialMessages={messages}
             initialTitle={title ?? ''}
+            initialFlyoutPositionMode={flyoutSettings.mode}
             onClose={() => {
+              setFlyoutSettings((prev) => ({ ...prev, isOpen: false }));
               setIsOpen(false);
             }}
+            onFlyoutPositionModeChange={(next) => {
+              setFlyoutSettings((prev) => ({ ...prev, mode: next }));
+            }}
             navigateToConversation={(conversationId?: string) => {
               application.navigateToUrl(
                 http.basePath.prepend(

x-pack/solutions/observability/plugins/observability_ai_assistant_app/public/hooks/use_local_storage.test.ts

@@ -10,7 +10,7 @@ import { useLocalStorage } from './use_local_storage';
 
 describe('useLocalStorage', () => {
   const key = 'testKey';
-  const defaultValue = 'defaultValue';
+  const defaultValue: string = 'defaultValue';
 
   beforeEach(() => {
     localStorage.clear();
@@ -44,17 +44,6 @@ describe('useLocalStorage', () => {
     expect(JSON.parse(localStorage.getItem(key) || '')).toBe(newValue);
   });
 
-  it('should remove the value from local storage when the value is undefined', () => {
-    const { result } = renderHook(() => useLocalStorage(key, defaultValue));
-    const [, saveToStorage] = result.current;
-
-    act(() => {
-      saveToStorage(undefined as unknown as string);
-    });
-
-    expect(localStorage.getItem(key)).toBe(null);
-  });
-
   it('should listen for storage events to window, and remove the listener upon unmount', () => {
     const addEventListenerSpy = jest.spyOn(window, 'addEventListener');
     const removeEventListenerSpy = jest.spyOn(window, 'removeEventListener');

x-pack/solutions/observability/plugins/observability_ai_assistant_app/public/hooks/use_local_storage.ts

@@ -5,45 +5,74 @@
  * 2.0.
  */
 
-import { useState, useEffect, useMemo, useCallback } from 'react';
+import { useEffect, useMemo, useRef, useState } from 'react';
 
-export function useLocalStorage<T>(key: string, defaultValue: T) {
-  // This is necessary to fix a race condition issue.
-  // It guarantees that the latest value will be always returned after the value is updated
-  const [storageUpdate, setStorageUpdate] = useState(0);
+const LOCAL_STORAGE_UPDATE_EVENT_TYPE = 'customLocalStorage';
 
-  const item = useMemo(() => {
-    return getFromStorage(key, defaultValue);
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [key, storageUpdate, defaultValue]);
-
-  const saveToStorage = useCallback(
-    (value: T) => {
-      if (value === undefined) {
-        window.localStorage.removeItem(key);
-      } else {
-        window.localStorage.setItem(key, JSON.stringify(value));
-        setStorageUpdate(storageUpdate + 1);
-      }
-    },
-    [key, storageUpdate]
+export function useLocalStorage<T extends AllowedValue>(
+  key: string,
+  defaultValue: T | (() => T)
+): [T, SetValue<T>] {
+  const defaultValueRef = useRef<T>(
+    typeof defaultValue === 'function' ? defaultValue() : defaultValue
   );
 
+  const [value, setValue] = useState(() => getFromStorage(key, defaultValueRef.current));
+
+  const valueRef = useRef(value);
+  valueRef.current = value;
+
+  const setter = useMemo(() => {
+    return (valueOrCallback: T | ((prev: T) => T)) => {
+      const nextValue =
+        typeof valueOrCallback === 'function' ? valueOrCallback(valueRef.current) : valueOrCallback;
+
+      window.localStorage.setItem(key, JSON.stringify(nextValue));
+
+      /*
+       * This is necessary to trigger the event listener in the same
+       * window context.
+       */
+      window.dispatchEvent(
+        new window.CustomEvent<{ key: string }>(LOCAL_STORAGE_UPDATE_EVENT_TYPE, {
+          detail: { key },
+        })
+      );
+    };
+  }, [key]);
+
   useEffect(() => {
-    function onUpdate(event: StorageEvent) {
+    function updateValueFromStorage() {
+      setValue(getFromStorage(key, defaultValueRef.current));
+    }
+
+    function onStorageEvent(event: StorageEvent) {
       if (event.key === key) {
-        setStorageUpdate(storageUpdate + 1);
+        updateValueFromStorage();
       }
     }
-    window.addEventListener('storage', onUpdate);
-    return () => {
-      window.removeEventListener('storage', onUpdate);
-    };
-  }, [key, setStorageUpdate, storageUpdate]);
 
-  return useMemo(() => [item, saveToStorage] as const, [item, saveToStorage]);
+    function onCustomLocalStorageEvent(event: Event) {
+      if (event instanceof window.CustomEvent && event.detail?.key === key) {
+        updateValueFromStorage();
+      }
+    }
+
+    window.addEventListener('storage', onStorageEvent);
+    window.addEventListener(LOCAL_STORAGE_UPDATE_EVENT_TYPE, onCustomLocalStorageEvent);
+
+    return () => {
+      window.removeEventListener('storage', onStorageEvent);
+      window.removeEventListener(LOCAL_STORAGE_UPDATE_EVENT_TYPE, onCustomLocalStorageEvent);
+    };
+  }, [key]);
+
+  return [value, setter];
 }
 
+type AllowedValue = string | number | boolean | Record<string, any> | any[];
+type SetValue<T extends AllowedValue> = (next: T | ((prev: T) => T)) => void;
+
 function getFromStorage<T>(keyName: string, defaultValue: T) {
   const storedItem = window.localStorage.getItem(keyName);