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

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

@@ -3,6 +3,8 @@ 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 { Chat } from './@ai-sdk/react/chat-react';
+import sinon from 'sinon';
 
 describe('AssistantChat', function () {
   const mockMessages: UIMessage[] = [
@@ -23,8 +25,30 @@ describe('AssistantChat', function () {
     },
   ];
 
+  let renderWithChat: (messages: UIMessage[]) => {
+    result: ReturnType<typeof render>;
+    chat: Chat<UIMessage> & {
+      sendMessage: sinon.SinonStub;
+    };
+  };
+
+  beforeEach(() => {
+    renderWithChat = (messages: UIMessage[]) => {
+      const newChat = new Chat<UIMessage>({
+        messages,
+      });
+      sinon.replace(newChat, 'sendMessage', sinon.stub());
+      return {
+        result: render(<AssistantChat chat={newChat} />),
+        chat: newChat as unknown as Chat<UIMessage> & {
+          sendMessage: sinon.SinonStub;
+        },
+      };
+    };
+  });
+
   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 +58,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 +71,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 +82,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 +96,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,7 +110,7 @@ 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;
@@ -95,27 +119,20 @@ describe('AssistantChat', function () {
       .exist;
   });
 
-  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 +147,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 +204,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 +223,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 (

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

@@ -1,19 +1,25 @@
-import { DrawerSection } from '@mongodb-js/compass-components';
-import React, { type PropsWithChildren, useCallback, useRef } from 'react';
-import { type UIMessage, useChat } from './@ai-sdk/react/use-chat';
+import React, { type PropsWithChildren, useRef } from 'react';
+import { type UIMessage } from './@ai-sdk/react/use-chat';
 import type { Chat } from './@ai-sdk/react/chat-react';
-import { AssistantChat } from './assistant-chat';
 import { usePreference } from 'compass-preferences-model/provider';
+import { createContext, useContext } from 'react';
 
 export const ASSISTANT_DRAWER_ID = 'compass-assistant-drawer';
 
-import { createContext, useContext } from 'react';
+interface AssistantContextType {
+  chat: Chat<UIMessage>;
+  isEnabled: boolean;
+}
 
-type AssistantActions = unknown;
+export const AssistantContext = createContext<AssistantContextType | null>(
+  null
+);
 
-export const AssistantActionsContext = createContext<AssistantActions>({});
+type AssistantActionsContextType = unknown;
+export const AssistantActionsContext =
+  createContext<AssistantActionsContextType>({});
 
-export function useAssistantActions(): AssistantActions {
+export function useAssistantActions(): AssistantActionsContextType {
   return useContext(AssistantActionsContext);
 }
 
@@ -24,35 +30,27 @@ export const AssistantProvider: React.FunctionComponent<
 > = ({ chat, children }) => {
   const enableAIAssistant = usePreference('enableAIAssistant');
 
-  const { messages, sendMessage } = useChat({
+  const assistantContext = useRef<AssistantContextType>({
     chat,
+    isEnabled: enableAIAssistant,
   });
+  assistantContext.current = {
+    chat,
+    isEnabled: enableAIAssistant,
+  };
 
-  const contextActions = useRef({});
-
-  const handleMessageSend = useCallback(
-    (messageBody: string) => {
-      void sendMessage({ text: messageBody });
-    },
-    [sendMessage]
-  );
+  const assistantActionsContext = useRef<AssistantActionsContextType>({});
 
   if (!enableAIAssistant) {
     return <>{children}</>;
   }
 
   return (
-    <AssistantActionsContext.Provider value={contextActions.current}>
-      <DrawerSection
-        id={ASSISTANT_DRAWER_ID}
-        title="MongoDB Assistant"
-        label="MongoDB Assistant"
-        glyph="Sparkle"
-      >
-        <AssistantChat messages={messages} onSendMessage={handleMessageSend} />
-      </DrawerSection>
-      {children}
-    </AssistantActionsContext.Provider>
+    <AssistantContext.Provider value={assistantContext.current}>
+      <AssistantActionsContext.Provider value={assistantActionsContext.current}>
+        {children}
+      </AssistantActionsContext.Provider>
+    </AssistantContext.Provider>
   );
 };
 

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

@@ -0,0 +1,36 @@
+import React, { useContext } from 'react';
+import { DrawerSection } from '@mongodb-js/compass-components';
+import { AssistantChat } from './assistant-chat';
+import { ASSISTANT_DRAWER_ID, AssistantContext } from './assistant-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 = () => {
+  const context = useContext(AssistantContext);
+
+  if (!context) {
+    throw new Error(
+      'CompassAssistantDrawer must be used within an CompassAssistantProvider'
+    );
+  }
+
+  if (!context.isEnabled) {
+    return null;
+  }
+
+  const { chat } = context;
+
+  return (
+    <DrawerSection
+      id={ASSISTANT_DRAWER_ID}
+      title="MongoDB Assistant"
+      label="MongoDB Assistant"
+      glyph="Sparkle"
+    >
+      <AssistantChat chat={chat} />
+    </DrawerSection>
+  );
+};

packages/compass-assistant/src/index.tsx

@@ -7,7 +7,7 @@ import type { PropsWithChildren } from 'react';
 import type { UIMessage } from 'ai';
 import React from 'react';
 
-const CompassAssistantProvider = registerCompassPlugin(
+export const CompassAssistantProvider = registerCompassPlugin(
   {
     name: 'CompassAssistant',
     component: ({
@@ -38,7 +38,4 @@ const CompassAssistantProvider = registerCompassPlugin(
   }
 );
 
-export { CompassAssistantProvider };
-
-// Export hooks and components for external use
-export { AssistantProvider, useAssistantActions } from './assistant-provider';
+export { CompassAssistantDrawer } from './compass-assistant-drawer';

packages/compass-web/src/entrypoint.tsx

@@ -58,7 +58,10 @@ import { WebWorkspaceTab as WelcomeWorkspaceTab } from '@mongodb-js/compass-welc
 import { useCompassWebPreferences } from './preferences';
 import { DataModelingWorkspaceTab as DataModelingWorkspace } from '@mongodb-js/compass-data-modeling';
 import { DataModelStorageServiceProviderInMemory } from '@mongodb-js/compass-data-modeling/web';
-import { CompassAssistantProvider } from '@mongodb-js/compass-assistant';
+import {
+  CompassAssistantDrawer,
+  CompassAssistantProvider,
+} from '@mongodb-js/compass-assistant';
 
 export type TrackFunction = (
   event: string,
@@ -219,6 +222,7 @@ function CompassWorkspace({
                   <CreateNamespacePlugin></CreateNamespacePlugin>
                   <DropNamespacePlugin></DropNamespacePlugin>
                   <RenameCollectionPlugin></RenameCollectionPlugin>
+                  <CompassAssistantDrawer />
                 </>
               );
             }}