feat(compass-collection): Add Disclaimer Screen - Mock Data Generator CLOUDP-333851

packages/compass-collection/package.json

@@ -52,6 +52,7 @@
     "@mongodb-js/compass-app-stores": "^7.57.0",
     "@mongodb-js/compass-components": "^1.49.0",
     "@mongodb-js/compass-connections": "^1.71.0",
+    "@mongodb-js/compass-generative-ai": "^0.51.0",
     "@mongodb-js/compass-logging": "^1.7.12",
     "@mongodb-js/compass-telemetry": "^1.14.0",
     "@mongodb-js/compass-workspaces": "^0.52.0",

packages/compass-collection/src/components/collection-header/collection-header.spec.tsx

@@ -26,7 +26,7 @@ function renderCollectionHeader(
   const mockStore = createStore(() => ({
     mockDataGenerator: {
       isModalOpen: false,
-      currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
+      currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
     },
   }));
 

packages/compass-collection/src/components/collection-header/collection-header.tsx

@@ -20,7 +20,7 @@ import { useConnectionInfo } from '@mongodb-js/compass-connections/provider';
 import { getConnectionTitle } from '@mongodb-js/connection-info';
 import MockDataGeneratorModal from '../mock-data-generator-modal/mock-data-generator-modal';
 import { connect } from 'react-redux';
-import { mockDataGeneratorModalOpened } from '../../modules/collection-tab';
+import { openMockDataGeneratorModal } from '../../modules/collection-tab';
 
 const collectionHeaderStyles = css({
   padding: spacing[400],
@@ -182,7 +182,7 @@ const CollectionHeader: React.FunctionComponent<CollectionHeaderProps> = ({
 };
 
 const ConnectedCollectionHeader = connect(undefined, {
-  onOpenMockDataModal: mockDataGeneratorModalOpened,
+  onOpenMockDataModal: openMockDataGeneratorModal,
 })(CollectionHeader);
 
 export default ConnectedCollectionHeader;

packages/compass-collection/src/components/mock-data-generator-modal/constants.ts

@@ -1,7 +1,6 @@
 import { MockDataGeneratorStep } from './types';
 
 export const StepButtonLabelMap = {
-  [MockDataGeneratorStep.AI_DISCLAIMER]: 'Use Natural Language',
   [MockDataGeneratorStep.SCHEMA_CONFIRMATION]: 'Confirm',
   [MockDataGeneratorStep.SCHEMA_EDITOR]: 'Next',
   [MockDataGeneratorStep.DOCUMENT_COUNT]: 'Next',

packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx

@@ -23,7 +23,7 @@ describe('MockDataGeneratorModal', () => {
 
   function renderModal({
     isOpen = true,
-    currentStep = MockDataGeneratorStep.AI_DISCLAIMER,
+    currentStep = MockDataGeneratorStep.SCHEMA_CONFIRMATION,
   } = {}) {
     return render(
       <MockDataGeneratorModal

packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.tsx

@@ -70,7 +70,7 @@ const MockDataGeneratorModal = ({
       <ModalFooter className={footerStyles}>
         <Button
           onClick={onPreviousStep}
-          disabled={currentStep === MockDataGeneratorStep.AI_DISCLAIMER}
+          disabled={currentStep === MockDataGeneratorStep.SCHEMA_CONFIRMATION}
         >
           Back
         </Button>

packages/compass-collection/src/components/mock-data-generator-modal/types.ts

@@ -1,5 +1,4 @@
 export enum MockDataGeneratorStep {
-  AI_DISCLAIMER = 'AI_DISCLAIMER',
   SCHEMA_CONFIRMATION = 'SCHEMA_CONFIRMATION',
   SCHEMA_EDITOR = 'SCHEMA_EDITOR',
   DOCUMENT_COUNT = 'DOCUMENT_COUNT',

packages/compass-collection/src/index.ts

@@ -18,6 +18,7 @@ import {
   CollectionWorkspaceTitle,
   CollectionPluginTitleComponent,
 } from './plugin-tab-title';
+import { atlasAiServiceLocator } from '@mongodb-js/compass-generative-ai/provider';
 
 export const WorkspaceTab: WorkspacePlugin<typeof CollectionWorkspaceTitle> = {
   name: CollectionWorkspaceTitle,
@@ -37,6 +38,7 @@ export const WorkspaceTab: WorkspacePlugin<typeof CollectionWorkspaceTitle> = {
       connectionInfoRef: connectionInfoRefLocator,
       logger: createLoggerLocator('COMPASS-COLLECTION'),
       preferences: preferencesLocator,
+      atlasAiService: atlasAiServiceLocator,
     }
   ),
   content: CollectionTab,

packages/compass-collection/src/modules/collection-tab.ts

@@ -10,6 +10,8 @@ import type { DataService } from '@mongodb-js/compass-connections/provider';
 import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
 import { type Logger, mongoLogId } from '@mongodb-js/compass-logging/provider';
 import { type PreferencesAccess } from 'compass-preferences-model/provider';
+import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider';
+
 import { isInternalFieldPath } from 'hadron-document';
 import toNS from 'mongodb-ns';
 import {
@@ -24,11 +26,11 @@ import {
 import { calculateSchemaDepth } from '../calculate-schema-depth';
 import { processSchema } from '../transform-schema-to-field-info';
 import type { Document, MongoError } from 'mongodb';
+import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
 
 const DEFAULT_SAMPLE_SIZE = 100;
 
 const NO_DOCUMENTS_ERROR = 'No documents found in the collection to analyze.';
-import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
 
 function isAction<A extends AnyAction>(
   action: AnyAction,
@@ -65,6 +67,7 @@ type CollectionThunkAction<R, A extends AnyAction = AnyAction> = ThunkAction<
     experimentationServices: ReturnType<typeof experimentationServiceLocator>;
     logger: Logger;
     preferences: PreferencesAccess;
+    atlasAiService: AtlasAiService;
   },
   A
 >;
@@ -148,7 +151,7 @@ const reducer: Reducer<CollectionState, Action> = (
     },
     mockDataGenerator: {
       isModalOpen: false,
-      currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
+      currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
     },
   },
   action
@@ -236,7 +239,7 @@ const reducer: Reducer<CollectionState, Action> = (
       mockDataGenerator: {
         ...state.mockDataGenerator,
         isModalOpen: true,
-        currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
+        currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
       },
     };
   }
@@ -266,9 +269,6 @@ const reducer: Reducer<CollectionState, Action> = (
     let nextStep: MockDataGeneratorStep;
 
     switch (currentStep) {
-      case MockDataGeneratorStep.AI_DISCLAIMER:
-        nextStep = MockDataGeneratorStep.SCHEMA_CONFIRMATION;
-        break;
       case MockDataGeneratorStep.SCHEMA_CONFIRMATION:
         nextStep = MockDataGeneratorStep.SCHEMA_EDITOR;
         break;
@@ -305,7 +305,8 @@ const reducer: Reducer<CollectionState, Action> = (
 
     switch (currentStep) {
       case MockDataGeneratorStep.SCHEMA_CONFIRMATION:
-        previousStep = MockDataGeneratorStep.AI_DISCLAIMER;
+        // TODO: Decide with product what we want behavior to be: close modal? Re-open disclaimer modal, if possible?
+        previousStep = MockDataGeneratorStep.SCHEMA_CONFIRMATION;
         break;
       case MockDataGeneratorStep.SCHEMA_EDITOR:
         previousStep = MockDataGeneratorStep.SCHEMA_CONFIRMATION;
@@ -372,6 +373,27 @@ export const selectTab = (
   };
 };
 
+export const openMockDataGeneratorModal = (): CollectionThunkAction<
+  Promise<void>
+> => {
+  return async (dispatch, _getState, { atlasAiService, logger }) => {
+    try {
+      if (process.env.COMPASS_E2E_SKIP_ATLAS_SIGNIN !== 'true') {
+        await atlasAiService.ensureAiFeatureAccess();
+      }
+      dispatch(mockDataGeneratorModalOpened());
+    } catch (error) {
+      // if failed or user canceled we just don't show the modal
+      logger.log.error(
+        mongoLogId(1_001_000_364),
+        'Collections',
+        'Failed to ensure AI feature access and open mock data generator modal',
+        error
+      );
+    }
+  };
+};
+
 export const analyzeCollectionSchema = (): CollectionThunkAction<
   Promise<void>
 > => {

packages/compass-collection/src/stores/collection-tab.spec.ts

@@ -63,6 +63,7 @@ describe('Collection Tab Content store', function () {
     .stub(collectionTabModule, 'analyzeCollectionSchema')
     .returns(async () => {});
   const dataService = {} as any;
+  const atlasAiService = {} as any;
   let store: ReturnType<typeof activatePlugin>['store'];
   let deactivate: ReturnType<typeof activatePlugin>['deactivate'];
 
@@ -106,6 +107,7 @@ describe('Collection Tab Content store', function () {
         connectionInfoRef: connectionInfoRef as any,
         logger,
         preferences,
+        atlasAiService,
       },
       { on() {}, cleanup() {} } as any
     ));

packages/compass-collection/src/stores/collection-tab.ts

@@ -22,6 +22,7 @@ import {
 } from 'compass-preferences-model/provider';
 import { ExperimentTestName } from '@mongodb-js/compass-telemetry/provider';
 import { SCHEMA_ANALYSIS_STATE_INITIAL } from '../schema-analysis-types';
+import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider';
 
 export type CollectionTabOptions = {
   /**
@@ -48,6 +49,7 @@ export type CollectionTabServices = {
   connectionInfoRef: ReturnType<typeof connectionInfoRefLocator>;
   logger: Logger;
   preferences: PreferencesAccess;
+  atlasAiService: AtlasAiService;
 };
 
 export function activatePlugin(
@@ -67,6 +69,7 @@ export function activatePlugin(
     connectionInfoRef,
     logger,
     preferences,
+    atlasAiService,
   } = services;
 
   if (!collectionModel) {
@@ -87,7 +90,7 @@ export function activatePlugin(
       },
       mockDataGenerator: {
         isModalOpen: false,
-        currentStep: MockDataGeneratorStep.AI_DISCLAIMER,
+        currentStep: MockDataGeneratorStep.SCHEMA_CONFIRMATION,
       },
     },
     applyMiddleware(
@@ -98,6 +101,7 @@ export function activatePlugin(
         experimentationServices,
         logger,
         preferences,
+        atlasAiService,
       })
     )
   );

packages/compass-components/src/components/modals/marketing-modal.tsx

@@ -0,0 +1,19 @@
+import React from 'react';
+import {
+  Body,
+  MarketingModal as LeafyGreenMarketingModal,
+} from '../leafygreen';
+import { withStackedComponentStyles } from '../../hooks/use-stacked-component';
+
+function MarketingModal({
+  children,
+  ...props
+}: React.ComponentProps<typeof LeafyGreenMarketingModal>): React.ReactElement {
+  return (
+    <LeafyGreenMarketingModal {...props}>
+      <Body as="div">{children}</Body>
+    </LeafyGreenMarketingModal>
+  );
+}
+
+export default withStackedComponentStyles(MarketingModal);

packages/compass-components/src/index.ts

@@ -14,6 +14,7 @@ export {
   cache,
 } from '@leafygreen-ui/emotion';
 import ConfirmationModal from './components/modals/confirmation-modal';
+import MarketingModal from './components/modals/marketing-modal';
 import type {
   ElectronFileDialogOptions,
   ElectronShowFileDialogProvider,
@@ -131,6 +132,7 @@ export {
   defaultSidebarWidth,
   createElectronFileInputBackend,
   createJSDomFileInputDummyBackend,
+  MarketingModal,
 };
 export {
   useFocusState,

packages/compass-generative-ai/src/atlas-ai-service.ts

@@ -278,14 +278,18 @@ export class AtlasAiService {
 
   async ensureAiFeatureAccess({ signal }: { signal?: AbortSignal } = {}) {
     if (this.preferences.getPreferences().enableUnauthenticatedGenAI) {
-      return getStore().dispatch(optIntoGenAIWithModalPrompt({ signal }));
+      return getStore().dispatch(
+        optIntoGenAIWithModalPrompt({ signal, isCloudOptIn: false })
+      );
     }
 
     // When the ai feature is attempted to be opened we make sure
     // the user is signed into Atlas and opted in.
 
     if (this.apiURLPreset === 'cloud') {
-      return getStore().dispatch(optIntoGenAIWithModalPrompt({ signal }));
+      return getStore().dispatch(
+        optIntoGenAIWithModalPrompt({ signal, isCloudOptIn: true })
+      );
     }
     return getStore().dispatch(signIntoAtlasWithModalPrompt({ signal }));
   }