feat(compass-assistant): ensure opt-in before sending the first message COMPASS-9759 (#7263)

* ensure opt-in before sending the first message

* don't use useAtlasAiServiceContext

* circular dep, inferred troubles

* don't include ensureOptInAndSend in useAssistantActions()

* also don't expose clearChat

* deps

* tests

* factor out common type

* function is optional

* fix the cached feature flag

* add options

* only open the drawer once

Author: lerouxb

package-lock.json

@@ -56,6 +56,7 @@
     "@mongodb-js/compass-telemetry": "^1.14.0",
     "@mongodb-js/connection-info": "^0.17.1",
     "@mongodb-js/compass-logging": "^1.7.12",
+    "@mongodb-js/compass-generative-ai": "^0.51.0",
     "mongodb-connection-string-url": "^3.0.1",
     "ai": "^5.0.26",
     "compass-preferences-model": "^2.51.0",

packages/compass-assistant/package.json

@@ -8,7 +8,11 @@ import {
 import { AssistantChat } from './assistant-chat';
 import { expect } from 'chai';
 import { createMockChat } from '../test/utils';
-import type { AssistantMessage } from './compass-assistant-provider';
+import {
+  AssistantActionsContext,
+  type AssistantMessage,
+} from './compass-assistant-provider';
+import sinon from 'sinon';
 
 describe('AssistantChat', function () {
   const mockMessages: AssistantMessage[] = [
@@ -38,10 +42,30 @@ describe('AssistantChat', function () {
     } = {}
   ) {
     const chat = createMockChat({ messages, status });
-    const result = render(<AssistantChat chat={chat} />);
+
+    // The chat component does not use chat.sendMessage() directly, it uses
+    // ensureOptInAndSend() via the AssistantActionsContext.
+    const ensureOptInAndSendStub = sinon
+      .stub()
+      .callsFake(async (message, options, callback) => {
+        // call the callback so we can test the tracking
+        callback();
+
+        await chat.sendMessage(message, options);
+      });
+
+    const assistantActionsContext = {
+      ensureOptInAndSend: ensureOptInAndSendStub,
+    };
+    const result = render(
+      <AssistantActionsContext.Provider value={assistantActionsContext as any}>
+        <AssistantChat chat={chat} />
+      </AssistantActionsContext.Provider>
+    );
     return {
       result,
       chat,
+      ensureOptInAndSendStub,
     };
   }
 
@@ -156,8 +180,8 @@ describe('AssistantChat', function () {
     );
   });
 
-  it('calls sendMessage when form is submitted', async function () {
-    const { chat, result } = renderWithChat([]);
+  it('calls ensureOptInAndSend when form is submitted', async function () {
+    const { ensureOptInAndSendStub, result } = renderWithChat([]);
     const { track } = result;
     const inputField = screen.getByPlaceholderText(
       'Ask MongoDB Assistant a question'
@@ -167,8 +191,7 @@ describe('AssistantChat', function () {
     userEvent.type(inputField, 'What is aggregation?');
     userEvent.click(sendButton);
 
-    expect(chat.sendMessage.calledWith({ text: 'What is aggregation?' })).to.be
-      .true;
+    expect(ensureOptInAndSendStub.called).to.be.true;
 
     await waitFor(() => {
       expect(track).to.have.been.calledWith('Assistant Prompt Submitted', {
@@ -193,7 +216,7 @@ describe('AssistantChat', function () {
   });
 
   it('trims whitespace from input before sending', async function () {
-    const { chat, result } = renderWithChat([]);
+    const { ensureOptInAndSendStub, result } = renderWithChat([]);
     const { track } = result;
 
     const inputField = screen.getByPlaceholderText(
@@ -203,8 +226,7 @@ describe('AssistantChat', function () {
     userEvent.type(inputField, '  What is sharding?  ');
     userEvent.click(screen.getByLabelText('Send message'));
 
-    expect(chat.sendMessage.calledWith({ text: 'What is sharding?' })).to.be
-      .true;
+    expect(ensureOptInAndSendStub.called).to.be.true;
 
     await waitFor(() => {
       expect(track).to.have.been.calledWith('Assistant Prompt Submitted', {
@@ -213,8 +235,8 @@ describe('AssistantChat', function () {
     });
   });
 
-  it('does not call sendMessage when input is empty or whitespace-only', function () {
-    const { chat } = renderWithChat([]);
+  it('does not call ensureOptInAndSend when input is empty or whitespace-only', function () {
+    const { ensureOptInAndSendStub } = renderWithChat([]);
 
     const inputField = screen.getByPlaceholderText(
       'Ask MongoDB Assistant a question'
@@ -223,12 +245,12 @@ describe('AssistantChat', function () {
 
     // Test empty input
     userEvent.click(chatForm);
-    expect(chat.sendMessage.notCalled).to.be.true;
+    expect(ensureOptInAndSendStub.notCalled).to.be.true;
 
     // Test whitespace-only input
     userEvent.type(inputField, '   ');
     userEvent.click(chatForm);
-    expect(chat.sendMessage.notCalled).to.be.true;
+    expect(ensureOptInAndSendStub.notCalled).to.be.true;
   });
 
   it('displays user and assistant messages with different styling', function () {

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

@@ -1,5 +1,6 @@
-import React, { useCallback } from 'react';
+import React, { useCallback, useContext } from 'react';
 import type { AssistantMessage } from './compass-assistant-provider';
+import { AssistantActionsContext } from './compass-assistant-provider';
 import type { Chat } from './@ai-sdk/react/chat-react';
 import { useChat } from './@ai-sdk/react/use-chat';
 import {
@@ -129,7 +130,8 @@ export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
 }) => {
   const track = useTelemetry();
   const darkMode = useDarkMode();
-  const { messages, sendMessage, status, error, clearError } = useChat({
+  const { ensureOptInAndSend } = useContext(AssistantActionsContext);
+  const { messages, status, error, clearError } = useChat({
     chat,
     onError: (error) => {
       track('Assistant Response Failed', () => ({
@@ -157,13 +159,14 @@ export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
     (messageBody: string) => {
       const trimmedMessageBody = messageBody.trim();
       if (trimmedMessageBody) {
-        track('Assistant Prompt Submitted', {
-          user_input_length: trimmedMessageBody.length,
+        void ensureOptInAndSend?.({ text: trimmedMessageBody }, {}, () => {
+          track('Assistant Prompt Submitted', {
+            user_input_length: trimmedMessageBody.length,
+          });
         });
-        void sendMessage({ text: trimmedMessageBody });
       }
     },
-    [sendMessage, track]
+    [track, ensureOptInAndSend]
   );
 
   const handleFeedback = useCallback(