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

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

@@ -1,5 +1,6 @@
 import React, { useCallback } from 'react';
 import type { AssistantMessage } from './compass-assistant-provider';
+import { useAssistantActions } from './compass-assistant-provider';
 import type { Chat } from './@ai-sdk/react/chat-react';
 import { useChat } from './@ai-sdk/react/use-chat';
 import {
@@ -110,7 +111,8 @@ export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
 }) => {
   const track = useTelemetry();
   const darkMode = useDarkMode();
-  const { messages, sendMessage, status, error, clearError } = useChat({
+  const { ensureOptInAndSend } = useAssistantActions();
+  const { messages, status, error, clearError } = useChat({
     chat,
     onError: (error) => {
       track('Assistant Response Failed', () => ({
@@ -136,13 +138,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(

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

@@ -20,6 +20,7 @@ import { usePreference } from 'compass-preferences-model/provider';
 import { createLoggerLocator } from '@mongodb-js/compass-logging/provider';
 import type { ConnectionInfo } from '@mongodb-js/connection-info';
 import { useTelemetry } from '@mongodb-js/compass-telemetry/provider';
+import { useAtlasAiServiceContext } from '@mongodb-js/compass-generative-ai/provider';
 
 export const ASSISTANT_DRAWER_ID = 'compass-assistant-drawer';
 
@@ -36,6 +37,9 @@ export const AssistantContext = createContext<AssistantContextType | null>(
   null
 );
 
+type SendMessage = Parameters<Chat<AssistantMessage>['sendMessage']>[0];
+type SendOptions = Parameters<Chat<AssistantMessage>['sendMessage']>[1];
+
 type AssistantActionsContextType = {
   interpretExplainPlan: ({
     namespace,
@@ -53,13 +57,19 @@ type AssistantActionsContextType = {
   }) => void;
   clearChat: () => void;
   tellMoreAboutInsight: (context: ProactiveInsightsContext) => void;
+  ensureOptInAndSend: (
+    message: SendMessage,
+    options: SendOptions,
+    callback: () => void
+  ) => Promise<void>;
 };
 export const AssistantActionsContext =
   createContext<AssistantActionsContextType>({
     interpretExplainPlan: () => {},
     interpretConnectionError: () => {},
     tellMoreAboutInsight: () => {},
     clearChat: () => {},
+    ensureOptInAndSend: async () => {},
   });
 
 export function useAssistantActions(): AssistantActionsContextType & {
@@ -98,6 +108,7 @@ export const AssistantProvider: React.FunctionComponent<
 > = ({ chat, children }) => {
   const { openDrawer } = useDrawerActions();
   const track = useTelemetry();
+  const atlasAiService = useAtlasAiServiceContext();
   const createEntryPointHandler = useRef(function <T>(
     entryPointName:
       | 'explain plan'
@@ -106,12 +117,18 @@ export const AssistantProvider: React.FunctionComponent<
     builder: (props: T) => EntryPointMessage
   ) {
     return (props: T) => {
-      openDrawer(ASSISTANT_DRAWER_ID);
       const { prompt, displayText } = builder(props);
-      void chat.sendMessage({ text: prompt, metadata: { displayText } }, {});
-      track('Assistant Entry Point Used', {
-        source: entryPointName,
-      });
+      void assistantActionsContext.current.ensureOptInAndSend(
+        { text: prompt, metadata: { displayText } },
+        {},
+        () => {
+          openDrawer(ASSISTANT_DRAWER_ID);
+
+          track('Assistant Entry Point Used', {
+            source: entryPointName,
+          });
+        }
+      );
     };
   });
   const assistantActionsContext = useRef<AssistantActionsContextType>({
@@ -130,6 +147,24 @@ export const AssistantProvider: React.FunctionComponent<
     clearChat: () => {
       chat.messages = [];
     },
+    ensureOptInAndSend: async (
+      message: SendMessage,
+      options: SendOptions,
+      callback: () => void
+    ) => {
+      try {
+        await atlasAiService.ensureAiFeatureAccess();
+      } catch {
+        // opt-in failed: just do nothing
+        return;
+      }
+
+      // Call the callback to indicate that the opt-in was successful. A good
+      // place to do tracking.
+      callback();
+
+      await chat.sendMessage(message);
+    },
   });
 
   return (

packages/compass-generative-ai/src/provider.tsx

@@ -36,7 +36,7 @@ export const AtlasAiServiceProvider: React.FC<{
   );
 });
 
-function useAtlasAiServiceContext(): AtlasAiService {
+export function useAtlasAiServiceContext(): AtlasAiService {
   const service = useContext(AtlasAiServiceContext);
   if (!service) {
     throw new Error('No AtlasAiService available in this context');