[Observability AI Assistant] Make sure insight > chat flyout works pr…

…operly

x-pack/plugins/observability_ai_assistant/public/components/action_menu_item/action_menu_item.tsx

@@ -34,11 +34,12 @@ export function ObservabilityAIAssistantActionMenuItem() {
 
   const [conversationId, setConversationId] = useState<string>();
 
-  const { conversation, displayedMessages, setDisplayedMessages, save } = useConversation({
-    conversationId,
-    connectorId: connectors.selectedConnector,
-    chatService: chatService.value,
-  });
+  const { conversation, displayedMessages, setDisplayedMessages, save, saveTitle } =
+    useConversation({
+      conversationId,
+      connectorId: connectors.selectedConnector,
+      chatService: chatService.value,
+    });
 
   if (!service.isEnabled()) {
     return null;
@@ -87,6 +88,9 @@ export function ObservabilityAIAssistantActionMenuItem() {
             onChatUpdate={(nextMessages) => {
               setDisplayedMessages(nextMessages);
             }}
+            onChatTitleSave={(newTitle) => {
+              saveTitle(newTitle);
+            }}
           />
         </ObservabilityAIAssistantChatServiceProvider>
       ) : null}

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.stories.tsx

@@ -32,6 +32,9 @@ const defaultProps: ChatFlyoutProps = {
   title: 'How is this working',
   messages: [getAssistantSetupMessage({ contexts: [] })],
   onClose: () => {},
+  onChatComplete: () => {},
+  onChatTitleSave: () => {},
+  onChatUpdate: () => {},
 };
 
 export const ChatFlyout = Template.bind({});

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.tsx

@@ -9,13 +9,10 @@ import { css } from '@emotion/css';
 import { i18n } from '@kbn/i18n';
 import React from 'react';
 import type { Message } from '../../../common/types';
-import { useAbortableAsync } from '../../hooks/use_abortable_async';
-import { useConversation } from '../../hooks/use_conversation';
 import { useCurrentUser } from '../../hooks/use_current_user';
 import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
 import { useKibana } from '../../hooks/use_kibana';
 import { useKnowledgeBase } from '../../hooks/use_knowledge_base';
-import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
 import { useObservabilityAIAssistantRouter } from '../../hooks/use_observability_ai_assistant_router';
 import { getConnectorsManagementHref } from '../../utils/get_connectors_management_href';
 import { ChatBody } from './chat_body';
@@ -36,14 +33,16 @@ export function ChatFlyout({
   onClose,
   onChatUpdate,
   onChatComplete,
+  onChatTitleSave,
 }: {
   title: string;
   messages: Message[];
   conversationId?: string;
   isOpen: boolean;
   onClose: () => void;
-  onChatUpdate?: (messages: Message[]) => void;
-  onChatComplete?: (messages: Message[]) => void;
+  onChatUpdate: (messages: Message[]) => void;
+  onChatComplete: (messages: Message[]) => void;
+  onChatTitleSave: (title: string) => void;
 }) {
   const { euiTheme } = useEuiTheme();
   const {
@@ -54,25 +53,10 @@ export function ChatFlyout({
 
   const connectors = useGenAIConnectors();
 
-  const service = useObservabilityAIAssistant();
-
-  const chatService = useAbortableAsync(
-    ({ signal }) => {
-      return service.start({ signal });
-    },
-    [service]
-  );
-
   const router = useObservabilityAIAssistantRouter();
 
   const knowledgeBase = useKnowledgeBase();
 
-  const { saveTitle } = useConversation({
-    conversationId,
-    chatService: chatService.value,
-    connectorId: connectors.selectedConnector,
-  });
-
   return isOpen ? (
     <EuiFlyout onClose={onClose}>
       <EuiFlexGroup
@@ -127,7 +111,7 @@ export function ChatFlyout({
               }
             }}
             onSaveTitle={(newTitle) => {
-              saveTitle(newTitle);
+              onChatTitleSave(newTitle);
             }}
           />
         </EuiFlexItem>

x-pack/plugins/observability_ai_assistant/public/components/insight/insight.tsx

@@ -7,33 +7,35 @@
 import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
 import { useKibana } from '@kbn/kibana-react-plugin/public';
 import { AbortError } from '@kbn/kibana-utils-plugin/common';
+import { last } from 'lodash';
 import React, { useCallback, useEffect, useMemo, useState } from 'react';
 import type { Subscription } from 'rxjs';
 import { MessageRole, type Message } from '../../../common/types';
+import { ObservabilityAIAssistantChatServiceProvider } from '../../context/observability_ai_assistant_chat_service_provider';
+import { useAbortableAsync } from '../../hooks/use_abortable_async';
+import { useConversation } from '../../hooks/use_conversation';
 import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
+import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
+import { useObservabilityAIAssistantChatService } from '../../hooks/use_observability_ai_assistant_chat_service';
 import type { PendingMessage } from '../../types';
+import { getConnectorsManagementHref } from '../../utils/get_connectors_management_href';
+import { RegenerateResponseButton } from '../buttons/regenerate_response_button';
+import { StartChatButton } from '../buttons/start_chat_button';
+import { StopGeneratingButton } from '../buttons/stop_generating_button';
 import { ChatFlyout } from '../chat/chat_flyout';
 import { ConnectorSelectorBase } from '../connector_selector/connector_selector_base';
 import { MessagePanel } from '../message_panel/message_panel';
 import { MessageText } from '../message_panel/message_text';
-import { RegenerateResponseButton } from '../buttons/regenerate_response_button';
-import { StartChatButton } from '../buttons/start_chat_button';
-import { StopGeneratingButton } from '../buttons/stop_generating_button';
-import { InsightBase } from './insight_base';
 import { MissingCredentialsCallout } from '../missing_credentials_callout';
-import { getConnectorsManagementHref } from '../../utils/get_connectors_management_href';
-import { useObservabilityAIAssistantChatService } from '../../hooks/use_observability_ai_assistant_chat_service';
-import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
-import { useAbortableAsync } from '../../hooks/use_abortable_async';
-import { ObservabilityAIAssistantChatServiceProvider } from '../../context/observability_ai_assistant_chat_service_provider';
+import { InsightBase } from './insight_base';
 
 function ChatContent({
   title,
-  messages,
+  initialMessages,
   connectorId,
 }: {
   title: string;
-  messages: Message[];
+  initialMessages: Message[];
   connectorId: string;
 }) {
   const chatService = useObservabilityAIAssistantChatService();
@@ -42,51 +44,84 @@ function ChatContent({
   const [loading, setLoading] = useState(false);
   const [subscription, setSubscription] = useState<Subscription | undefined>();
 
+  const [conversationId, setConversationId] = useState<string>();
+
+  const { displayedMessages, setDisplayedMessages, save, saveTitle } = useConversation({
+    conversationId,
+    connectorId,
+    chatService,
+  });
+
   const reloadReply = useCallback(() => {
     setLoading(true);
 
+    let lastPendingMessage: PendingMessage | undefined;
+
     const nextSubscription = chatService
-      .chat({ messages, connectorId, function: 'none' })
+      .chat({ messages: initialMessages, connectorId, function: 'none' })
       .subscribe({
         next: (msg) => {
+          lastPendingMessage = msg;
           setPendingMessage(() => msg);
         },
         complete: () => {
+          setDisplayedMessages((prevMessages) =>
+            prevMessages.concat({
+              '@timestamp': new Date().toISOString(),
+              message: {
+                ...lastPendingMessage!.message,
+              },
+            })
+          );
           setLoading(false);
         },
       });
 
     setSubscription(nextSubscription);
-  }, [messages, connectorId, chatService]);
+  }, [initialMessages, setDisplayedMessages, connectorId, chatService]);
 
   useEffect(() => {
     reloadReply();
   }, [reloadReply]);
 
+  useEffect(() => {
+    setDisplayedMessages(initialMessages);
+  }, [initialMessages, setDisplayedMessages]);
+
   const [isOpen, setIsOpen] = useState(false);
 
-  const displayedMessages = useMemo(() => {
+  const messagesWithPending = useMemo(() => {
     return pendingMessage
-      ? messages.concat({
+      ? displayedMessages.concat({
           '@timestamp': new Date().toISOString(),
           message: {
             ...pendingMessage.message,
           },
         })
-      : messages;
-  }, [pendingMessage, messages]);
+      : displayedMessages;
+  }, [pendingMessage, displayedMessages]);
+
+  const lastMessage = last(messagesWithPending);
 
   return (
     <>
       <MessagePanel
-        body={<MessageText content={pendingMessage?.message.content ?? ''} loading={loading} />}
+        body={<MessageText content={lastMessage?.message.content ?? ''} loading={loading} />}
         error={pendingMessage?.error}
         controls={
           loading ? (
             <StopGeneratingButton
               onClick={() => {
                 subscription?.unsubscribe();
                 setLoading(false);
+                setDisplayedMessages((prevMessages) =>
+                  prevMessages.concat({
+                    '@timestamp': new Date().toISOString(),
+                    message: {
+                      ...pendingMessage!.message,
+                    },
+                  })
+                );
                 setPendingMessage((prev) => ({
                   message: {
                     role: MessageRole.Assistant,
@@ -124,6 +159,20 @@ function ChatContent({
           setIsOpen(() => false);
         }}
         messages={displayedMessages}
+        conversationId={conversationId}
+        onChatComplete={(nextMessages) => {
+          save(nextMessages)
+            .then((nextConversation) => {
+              setConversationId(nextConversation.conversation.id);
+            })
+            .catch(() => {});
+        }}
+        onChatUpdate={(nextMessages) => {
+          setDisplayedMessages(nextMessages);
+        }}
+        onChatTitleSave={(newTitle) => {
+          saveTitle(newTitle);
+        }}
       />
     </>
   );
@@ -151,7 +200,11 @@ export function Insight({ messages, title }: { messages: Message[]; title: strin
 
   if (hasOpened && connectors.selectedConnector) {
     children = (
-      <ChatContent title={title} messages={messages} connectorId={connectors.selectedConnector} />
+      <ChatContent
+        title={title}
+        initialMessages={messages}
+        connectorId={connectors.selectedConnector}
+      />
     );
   } else if (!connectors.loading && !connectors.connectors?.length) {
     children = (