chore: add better LG loading state

Author: gagik

packages/compass-assistant/src/assistant-chat.spec.tsx

@@ -29,8 +29,15 @@ describe('AssistantChat', function () {
     },
   ];
 
-  function renderWithChat(messages: AssistantMessage[]) {
-    const chat = createMockChat({ messages });
+  function renderWithChat(
+    messages: AssistantMessage[],
+    {
+      status,
+    }: {
+      status?: 'submitted' | 'streaming';
+    } = {}
+  ) {
+    const chat = createMockChat({ messages, status });
     const result = render(<AssistantChat chat={chat} />);
     return {
       result,
@@ -81,6 +88,16 @@ describe('AssistantChat', function () {
       .exist;
   });
 
+  it('displays loading state when chat status is submitted', function () {
+    renderWithChat([], { status: 'submitted' });
+    expect(screen.getByText(/MongoDB Assistant is thinking/)).to.exist;
+  });
+
+  it('does not display loading in all other cases', function () {
+    renderWithChat(mockMessages, { status: 'streaming' });
+    expect(screen.queryByText(/MongoDB Assistant is thinking/)).to.not.exist;
+  });
+
   it('send button is disabled when input is empty', function () {
     renderWithChat([]);
 

packages/compass-assistant/src/assistant-chat.tsx

@@ -228,13 +228,6 @@ export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
                 )}
               </Message>
             ))}
-            {status === 'submitted' && (
-              <Message
-                id="loading"
-                messageBody="Thinking..."
-                isSender={false}
-              />
-            )}
           </MessageFeed>
           {error && (
             <div className={errorBannerWrapperStyles}>
@@ -254,6 +247,7 @@ export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
             data-testid="assistant-chat-input"
             onMessageSend={handleMessageSend}
             className={inputBarFixesStyles}
+            state={status === 'submitted' ? 'loading' : undefined}
             textareaProps={{
               placeholder: 'Ask MongoDB Assistant a question',
             }}

packages/compass-assistant/test/utils.tsx

@@ -4,13 +4,18 @@ import type { AssistantMessage } from '../src/compass-assistant-provider';
 
 export const createMockChat = ({
   messages,
+  status,
 }: {
   messages: AssistantMessage[];
+  status?: 'submitted' | 'streaming';
 }) => {
   const newChat = new Chat<AssistantMessage>({
     messages,
   });
   sinon.replace(newChat, 'sendMessage', sinon.stub());
+  if (status) {
+    sinon.replaceGetter(newChat, 'status', () => status);
+  }
   return newChat as unknown as Chat<AssistantMessage> & {
     sendMessage: sinon.SinonStub;
   };