chore(compass-assistant): split drawer and provider usage for connections context COMPASS-9603 (#7199)

Author: gagik

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

@@ -3,16 +3,17 @@ import { render, screen, userEvent } from '@mongodb-js/testing-library-compass';
 import { AssistantChat } from './assistant-chat';
 import { expect } from 'chai';
 import type { UIMessage } from './@ai-sdk/react/use-chat';
+import { createMockChat } from '../test/utils';
 
 describe('AssistantChat', function () {
   const mockMessages: UIMessage[] = [
     {
-      id: '1',
+      id: 'user',
       role: 'user',
       parts: [{ type: 'text', text: 'Hello, MongoDB Assistant!' }],
     },
     {
-      id: '2',
+      id: 'assistant',
       role: 'assistant',
       parts: [
         {
@@ -23,8 +24,16 @@ describe('AssistantChat', function () {
     },
   ];
 
+  function renderWithChat(messages: UIMessage[]) {
+    const chat = createMockChat({ messages });
+    return {
+      result: render(<AssistantChat chat={chat} />),
+      chat,
+    };
+  }
+
   it('renders input field and send button', function () {
-    render(<AssistantChat messages={[]} />);
+    renderWithChat([]);
 
     const inputField = screen.getByTestId('assistant-chat-input');
     const sendButton = screen.getByTestId('assistant-chat-send-button');
@@ -34,7 +43,7 @@ describe('AssistantChat', function () {
   });
 
   it('input field accepts text input', function () {
-    render(<AssistantChat messages={[]} />);
+    renderWithChat([]);
 
     // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
     const inputField = screen.getByTestId(
@@ -47,7 +56,7 @@ describe('AssistantChat', function () {
   });
 
   it('send button is disabled when input is empty', function () {
-    render(<AssistantChat messages={[]} />);
+    renderWithChat([]);
 
     // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
     const sendButton = screen.getByTestId(
@@ -58,7 +67,7 @@ describe('AssistantChat', function () {
   });
 
   it('send button is enabled when input has text', function () {
-    render(<AssistantChat messages={[]} />);
+    renderWithChat([]);
 
     const inputField = screen.getByTestId('assistant-chat-input');
     // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
@@ -72,7 +81,7 @@ describe('AssistantChat', function () {
   });
 
   it('send button is disabled for whitespace-only input', function () {
-    render(<AssistantChat messages={[]} />);
+    renderWithChat([]);
 
     const inputField = screen.getByTestId('assistant-chat-input');
     // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
@@ -86,36 +95,32 @@ describe('AssistantChat', function () {
   });
 
   it('displays messages in the chat feed', function () {
-    render(<AssistantChat messages={mockMessages} />);
+    renderWithChat(mockMessages);
 
     expect(screen.getByTestId('assistant-message-user')).to.exist;
     expect(screen.getByTestId('assistant-message-assistant')).to.exist;
-    expect(screen.getByText('Hello, MongoDB Assistant!')).to.exist;
-    expect(screen.getByText('Hello! How can I help you with MongoDB today?')).to
-      .exist;
+    expect(screen.getByTestId('assistant-message-user')).to.have.text(
+      'Hello, MongoDB Assistant!'
+    );
+    expect(screen.getByTestId('assistant-message-assistant')).to.have.text(
+      'Hello! How can I help you with MongoDB today?'
+    );
   });
 
-  it('calls onSendMessage when form is submitted', function () {
-    let sentMessage = '';
-    const handleSendMessage = (message: string) => {
-      sentMessage = message;
-    };
-
-    render(<AssistantChat messages={[]} onSendMessage={handleSendMessage} />);
-
+  it('calls sendMessage when form is submitted', function () {
+    const { chat } = renderWithChat([]);
     const inputField = screen.getByTestId('assistant-chat-input');
     const sendButton = screen.getByTestId('assistant-chat-send-button');
 
     userEvent.type(inputField, 'What is aggregation?');
     userEvent.click(sendButton);
 
-    expect(sentMessage).to.equal('What is aggregation?');
+    expect(chat.sendMessage.calledWith({ text: 'What is aggregation?' })).to.be
+      .true;
   });
 
   it('clears input field after successful submission', function () {
-    const handleSendMessage = () => {};
-
-    render(<AssistantChat messages={[]} onSendMessage={handleSendMessage} />);
+    renderWithChat([]);
 
     // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
     const inputField = screen.getByTestId(
@@ -130,44 +135,35 @@ describe('AssistantChat', function () {
   });
 
   it('trims whitespace from input before sending', function () {
-    let sentMessage = '';
-    const handleSendMessage = (message: string) => {
-      sentMessage = message;
-    };
-
-    render(<AssistantChat messages={[]} onSendMessage={handleSendMessage} />);
+    const { chat } = renderWithChat([]);
 
     const inputField = screen.getByTestId('assistant-chat-input');
 
     userEvent.type(inputField, '  What is sharding?  ');
     userEvent.click(screen.getByTestId('assistant-chat-send-button'));
 
-    expect(sentMessage).to.equal('What is sharding?');
+    expect(chat.sendMessage.calledWith({ text: 'What is sharding?' })).to.be
+      .true;
   });
 
-  it('does not call onSendMessage when input is empty or whitespace-only', function () {
-    let messageSent = false;
-    const handleSendMessage = () => {
-      messageSent = true;
-    };
-
-    render(<AssistantChat messages={[]} onSendMessage={handleSendMessage} />);
+  it('does not call sendMessage when input is empty or whitespace-only', function () {
+    const { chat } = renderWithChat([]);
 
     const inputField = screen.getByTestId('assistant-chat-input');
     const chatForm = screen.getByTestId('assistant-chat-form');
 
     // Test empty input
     userEvent.click(chatForm);
-    expect(messageSent).to.be.false;
+    expect(chat.sendMessage.notCalled).to.be.true;
 
     // Test whitespace-only input
     userEvent.type(inputField, '   ');
     userEvent.click(chatForm);
-    expect(messageSent).to.be.false;
+    expect(chat.sendMessage.notCalled).to.be.true;
   });
 
   it('displays user and assistant messages with different styling', function () {
-    render(<AssistantChat messages={mockMessages} />);
+    renderWithChat(mockMessages);
 
     const userMessage = screen.getByTestId('assistant-message-user');
     const assistantMessage = screen.getByTestId('assistant-message-assistant');
@@ -196,7 +192,7 @@ describe('AssistantChat', function () {
       },
     ];
 
-    render(<AssistantChat messages={messagesWithMultipleParts} />);
+    renderWithChat(messagesWithMultipleParts);
 
     expect(screen.getByText('Here is part 1. And here is part 2.')).to.exist;
   });
@@ -215,7 +211,7 @@ describe('AssistantChat', function () {
       },
     ];
 
-    render(<AssistantChat messages={messagesWithMixedParts} />);
+    renderWithChat(messagesWithMixedParts);
 
     expect(screen.getByText('This is text content. More text content.')).to
       .exist;

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

@@ -1,30 +1,33 @@
 import React, { useCallback, useState } from 'react';
 import type { UIMessage } from './@ai-sdk/react/use-chat';
+import type { Chat } from './@ai-sdk/react/chat-react';
+import { useChat } from './@ai-sdk/react/use-chat';
 
 interface AssistantChatProps {
-  messages: UIMessage[];
-  onSendMessage?: (message: string) => void;
+  chat: Chat<UIMessage>;
 }
 
 /**
  * This component is currently using placeholders as Leafygreen UI updates are not available yet.
  * Before release, we will replace this with the actual Leafygreen chat components.
  */
 export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
-  messages,
-  onSendMessage,
+  chat,
 }) => {
   const [inputValue, setInputValue] = useState('');
+  const { messages, sendMessage } = useChat({
+    chat,
+  });
 
   const handleInputSubmit = useCallback(
     (e: React.FormEvent) => {
       e.preventDefault();
-      if (inputValue.trim() && onSendMessage) {
-        onSendMessage(inputValue.trim());
+      if (inputValue.trim()) {
+        void sendMessage({ text: inputValue.trim() });
         setInputValue('');
       }
     },
-    [inputValue, onSendMessage]
+    [inputValue, sendMessage]
   );
 
   return (
@@ -53,7 +56,7 @@ export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
         {messages.map((message) => (
           <div
             key={message.id}
-            data-testid={`assistant-message-${message.role}`}
+            data-testid={`assistant-message-${message.id}`}
             style={{
               marginBottom: '12px',
               padding: '8px 12px',

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

@@ -0,0 +1,43 @@
+import React, { useContext } from 'react';
+import { DrawerSection } from '@mongodb-js/compass-components';
+import { AssistantChat } from './assistant-chat';
+import {
+  ASSISTANT_DRAWER_ID,
+  AssistantContext,
+} from './compass-assistant-provider';
+import { usePreference } from 'compass-preferences-model/provider';
+
+/**
+ * CompassAssistantDrawer component that wraps AssistantChat in a DrawerSection.
+ * This component can be placed at any level in the component tree as long as
+ * it's within an AssistantProvider.
+ */
+export const CompassAssistantDrawer: React.FunctionComponent<{
+  autoOpen?: boolean;
+}> = ({ autoOpen }) => {
+  const chat = useContext(AssistantContext);
+
+  const enableAIAssistant = usePreference('enableAIAssistant');
+
+  if (!enableAIAssistant) {
+    return null;
+  }
+
+  if (!chat) {
+    throw new Error(
+      'CompassAssistantDrawer must be used within an CompassAssistantProvider'
+    );
+  }
+
+  return (
+    <DrawerSection
+      id={ASSISTANT_DRAWER_ID}
+      title="MongoDB Assistant"
+      label="MongoDB Assistant"
+      glyph="Sparkle"
+      autoOpen={autoOpen}
+    >
+      <AssistantChat chat={chat} />
+    </DrawerSection>
+  );
+};