Sync new commits from agent framework to main

public/components/__tests__/chat_experimental_badge.test.tsx

@@ -0,0 +1,53 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React from 'react';
+import { act, fireEvent, render, waitFor } from '@testing-library/react';
+import { I18nProvider } from '@osd/i18n/react';
+
+import { ChatExperimentalBadge } from '../chat_experimental_badge';
+
+describe('<ChatWindowHeaderTitle />', () => {
+  it('should show experimental dropdown after icon clicked', () => {
+    const { getByRole, getByText, queryByText } = render(
+      <I18nProvider>
+        <ChatExperimentalBadge />
+      </I18nProvider>
+    );
+
+    expect(queryByText('Experimental')).not.toBeInTheDocument();
+    fireEvent.click(getByRole('button'));
+    expect(getByText('Experimental')).toBeInTheDocument();
+  });
+
+  it('should hide experimental dropdown after click other places', async () => {
+    const { getByRole, getByText, queryByText } = render(
+      <I18nProvider>
+        <ChatExperimentalBadge />
+      </I18nProvider>
+    );
+
+    act(() => {
+      fireEvent.click(getByRole('button'));
+    });
+
+    await waitFor(() => {
+      expect(getByText('Experimental')).toBeInTheDocument();
+
+      // Ensure focus trap enabled, then we can click outside.
+      expect(
+        getByText('Experimental').closest('div[data-focus-lock-disabled="false"]')
+      ).toBeInTheDocument();
+    });
+
+    act(() => {
+      fireEvent.mouseDown(document.body);
+    });
+
+    await waitFor(() => {
+      expect(queryByText('Experimental')).not.toBeInTheDocument();
+    });
+  });
+});

public/components/__tests__/chat_window_header_title.test.tsx

@@ -14,10 +14,14 @@ import * as useChatActionsExports from '../../hooks/use_chat_actions';
 import * as useSaveChatExports from '../../hooks/use_save_chat';
 import * as chatContextExports from '../../contexts/chat_context';
 import * as coreContextExports from '../../contexts/core_context';
+import { IMessage } from '../../../common/types/chat_saved_object_attributes';
 
 import { ChatWindowHeaderTitle } from '../chat_window_header_title';
 
-const setup = () => {
+const setup = ({
+  messages = [],
+  ...rest
+}: { messages?: IMessage[]; sessionId?: string | undefined } = {}) => {
   const useCoreMock = {
     services: {
       ...coreMock.createStart(),
@@ -38,10 +42,10 @@ const setup = () => {
   useCoreMock.services.http.put.mockImplementation(() => Promise.resolve());
 
   const useChatStateMock = {
-    chatState: { messages: [] },
+    chatState: { messages },
   };
   const useChatContextMock = {
-    sessionId: '1',
+    sessionId: 'sessionId' in rest ? rest.sessionId : '1',
     title: 'foo',
     setSessionId: jest.fn(),
     setTitle: jest.fn(),
@@ -68,6 +72,7 @@ const setup = () => {
     useCoreMock,
     useChatStateMock,
     useChatContextMock,
+    useChatActionsMock,
     renderResult,
   };
 };
@@ -100,4 +105,106 @@ describe('<ChatWindowHeaderTitle />', () => {
       expect(useCoreMock.services.sessions.reload).toHaveBeenCalled();
     });
   });
+
+  it('should show "Rename conversation", "New conversation" and "Save to notebook" actions after title click', () => {
+    const { renderResult } = setup();
+
+    expect(
+      renderResult.queryByRole('button', { name: 'Rename conversation' })
+    ).not.toBeInTheDocument();
+    expect(
+      renderResult.queryByRole('button', { name: 'New conversation' })
+    ).not.toBeInTheDocument();
+    expect(
+      renderResult.queryByRole('button', { name: 'Save to notebook' })
+    ).not.toBeInTheDocument();
+
+    act(() => {
+      fireEvent.click(renderResult.getByText('foo'));
+    });
+
+    expect(renderResult.getByRole('button', { name: 'Rename conversation' })).toBeInTheDocument();
+    expect(renderResult.getByRole('button', { name: 'New conversation' })).toBeInTheDocument();
+    expect(renderResult.getByRole('button', { name: 'Save to notebook' })).toBeInTheDocument();
+  });
+
+  it('should show rename modal and hide rename actions after rename button clicked', async () => {
+    const { renderResult } = setup();
+
+    act(() => {
+      fireEvent.click(renderResult.getByText('foo'));
+    });
+
+    act(() => {
+      fireEvent.click(renderResult.getByRole('button', { name: 'Rename conversation' }));
+    });
+
+    await waitFor(() => {
+      expect(renderResult.getByText('Edit conversation name')).toBeInTheDocument();
+      expect(
+        renderResult.queryByRole('button', { name: 'Rename conversation' })
+      ).not.toBeInTheDocument();
+    });
+  });
+
+  it('should call loadChat with undefined, hide actions and show success toasts after new conversation button clicked', async () => {
+    const { renderResult, useCoreMock, useChatActionsMock } = setup();
+
+    act(() => {
+      fireEvent.click(renderResult.getByText('foo'));
+    });
+
+    expect(useChatActionsMock.loadChat).not.toHaveBeenCalled();
+    expect(useCoreMock.services.notifications.toasts.addSuccess).not.toHaveBeenCalled();
+
+    act(() => {
+      fireEvent.click(renderResult.getByRole('button', { name: 'New conversation' }));
+    });
+
+    await waitFor(() => {
+      expect(useChatActionsMock.loadChat).toHaveBeenCalledWith(undefined);
+      expect(
+        renderResult.queryByRole('button', { name: 'New conversation' })
+      ).not.toBeInTheDocument();
+      expect(useCoreMock.services.notifications.toasts.addSuccess).toHaveBeenCalledWith(
+        'A new conversation is started and the previous one is saved.'
+      );
+    });
+  });
+
+  it('should show save to notebook modal after "Save to notebook" clicked', async () => {
+    const { renderResult } = setup();
+
+    act(() => {
+      fireEvent.click(renderResult.getByText('foo'));
+    });
+
+    act(() => {
+      fireEvent.click(renderResult.getByRole('button', { name: 'Save to notebook' }));
+    });
+
+    await waitFor(() => {
+      expect(renderResult.queryByText('Save to notebook')).toBeInTheDocument();
+    });
+  });
+
+  it('should disable "Save to notebook" button when message does not include input', async () => {
+    const { renderResult } = setup({
+      messages: [{ type: 'output', content: 'bar', contentType: 'markdown' }],
+    });
+
+    act(() => {
+      fireEvent.click(renderResult.getByText('foo'));
+    });
+
+    expect(renderResult.getByRole('button', { name: 'Save to notebook' })).toBeDisabled();
+  });
+
+  it('should show "OpenSearch Assistant" when sessionId is undefined', async () => {
+    const { renderResult } = setup({
+      sessionId: undefined,
+    });
+
+    expect(renderResult.getByText('OpenSearch Assistant')).toBeInTheDocument();
+  });
 });

public/components/chat_experimental_badge.tsx

@@ -31,7 +31,14 @@ export const ChatExperimentalBadge = ({ onClick }: ChatExperimentalBadgeProps) =
   return (
     <EuiPopover
       isOpen={visible}
-      button={<EuiButtonIcon color="text" iconType="beaker" onClick={handleIconClick} />}
+      button={
+        <EuiButtonIcon
+          color="text"
+          iconType="beaker"
+          onClick={handleIconClick}
+          aria-label="Experimental badge"
+        />
+      }
       closePopover={closePopover}
       onClick={onClick}
     >

public/components/chat_window_header_title.tsx

@@ -25,6 +25,7 @@ export const ChatWindowHeaderTitle = React.memo(() => {
   const core = useCore();
   const [isPopoverOpen, setPopoverOpen] = useState(false);
   const [isRenameModalOpen, setRenameModalOpen] = useState(false);
+  const [isSaveNotebookModalOpen, setSaveNotebookModalOpen] = useState(false);
   const { chatState } = useChatState();
   const { saveChat } = useSaveChat();
 
@@ -50,6 +51,10 @@ export const ChatWindowHeaderTitle = React.memo(() => {
     [chatContext, core.services.sessions]
   );
 
+  const handleSaveNotebookModalClose = () => {
+    setSaveNotebookModalOpen(false);
+  };
+
   const button = (
     <EuiFlexGroup
       style={{ maxWidth: '300px', padding: '0 8px' }}
@@ -100,10 +105,8 @@ export const ChatWindowHeaderTitle = React.memo(() => {
     <EuiContextMenuItem
       key="save-as-notebook"
       onClick={() => {
-        const modal = core.overlays.openModal(
-          <NotebookNameModal onClose={() => modal.close()} saveChat={saveChat} />
-        );
         closePopover();
+        setSaveNotebookModalOpen(true);
       }}
       // User only can save conversation when he send a message at least.
       disabled={chatState.messages.every((item) => item.type !== 'input')}
@@ -131,6 +134,9 @@ export const ChatWindowHeaderTitle = React.memo(() => {
           defaultTitle={chatContext.title!}
         />
       )}
+      {isSaveNotebookModalOpen && (
+        <NotebookNameModal onClose={handleSaveNotebookModalClose} saveChat={saveChat} />
+      )}
     </>
   );
 });