fix(preferences, settings, atlas-service): ensure Atlas Login doesn't show in settings if user didn't get the Gen AI feature rollout (#4912)

chore(preferences, settings, atlas-service): use the same check everywhere for whether or not the feature is enabled

Author: gribnoysup

packages/atlas-service/src/main.ts

@@ -41,7 +41,7 @@ import preferences from 'compass-preferences-model';
 import { SecretStore, SECRET_STORE_KEY } from './secret-store';
 import { AtlasUserConfigStore } from './user-config-store';
 import { OidcPluginLogger } from './oidc-plugin-logger';
-import { getActiveUser } from 'compass-preferences-model';
+import { getActiveUser, isAIFeatureEnabled } from 'compass-preferences-model';
 import { spawn } from 'child_process';
 
 const { log, track } = createLoggerAndTelemetry('COMPASS-ATLAS-SERVICE');
@@ -92,11 +92,7 @@ export async function throwIfNotOk(
 }
 
 function throwIfAINotEnabled(atlasService: typeof AtlasService) {
-  if (
-    (!preferences.getPreferences().cloudFeatureRolloutAccess?.GEN_AI_COMPASS &&
-      !preferences.getPreferences().enableAIWithoutRolloutAccess) ||
-    !preferences.getPreferences().enableAIFeatures
-  ) {
+  if (!isAIFeatureEnabled()) {
     throw new Error(
       "Compass' AI functionality is not currently enabled. Please try again later."
     );

packages/compass-e2e-tests/tests/collection-ai-query.test.ts

@@ -44,7 +44,7 @@ describe('Collection ai query', function () {
 
     telemetry = await startTelemetryServer();
     compass = await beforeTests({
-      extraSpawnArgs: ['--enableAIExperience'],
+      extraSpawnArgs: ['--enableGenAIExperience'],
     });
     browser = compass.browser;
   });

packages/compass-preferences-model/src/feature-flags.ts

@@ -14,7 +14,7 @@ export type FeatureFlagDefinition = {
 };
 
 export type FeatureFlags = {
-  enableAIExperience: boolean;
+  enableGenAIExperience: boolean;
   enableAIWithoutRolloutAccess: boolean;
   enableLgDarkmode: boolean;
   enableOidc: boolean; // Not capitalized "OIDC" for spawn arg casing.
@@ -31,7 +31,7 @@ export const featureFlags: Required<{
    * Feature flag for enabling the natural text input on the query bar.
    * Epic: COMPASS-6866
    */
-  enableAIExperience: {
+  enableGenAIExperience: {
     stage: 'released',
     description: {
       short: 'Compass AI Features',

packages/compass-preferences-model/src/index.ts

@@ -29,6 +29,7 @@ export {
   setupPreferencesAndUser,
   getActiveUser,
   useIsAIFeatureEnabled,
+  isAIFeatureEnabled,
 } from './utils';
 export type { User } from './storage';
 

packages/compass-preferences-model/src/preferences.spec.ts

@@ -151,7 +151,7 @@ describe('Preferences class', function () {
       enableDevTools: 'set-global',
       networkTraffic: 'set-global',
       trackUsageStatistics: 'set-global',
-      enableAIFeatures: 'set-global',
+      enableGenAIFeatures: 'set-global',
       enableMaps: 'set-cli',
       enableShell: 'set-cli',
       readOnly: 'set-global',
@@ -215,7 +215,7 @@ describe('Preferences class', function () {
       },
       {
         networkTraffic: false,
-        enableAIFeatures: false,
+        enableGenAIFeatures: false,
         enableMaps: false,
         enableFeedbackPanel: false,
         trackUsageStatistics: false,
@@ -248,7 +248,7 @@ describe('Preferences class', function () {
 
     const states = preferences.getPreferenceStates();
     expect(states).to.deep.equal({
-      enableAIFeatures: 'hardcoded',
+      enableGenAIFeatures: 'hardcoded',
       enableDevTools: 'set-global',
       enableMaps: 'set-cli',
       enableFeedbackPanel: 'hardcoded',

packages/compass-preferences-model/src/preferences.ts

@@ -26,7 +26,7 @@ export type UserConfigurablePreferences = PermanentFeatureFlags &
   FeatureFlags & {
     // User-facing preferences
     autoUpdates: boolean;
-    enableAIFeatures: boolean;
+    enableGenAIFeatures: boolean;
     enableMaps: boolean;
     trackUsageStatistics: boolean;
     enableFeedbackPanel: boolean;
@@ -439,15 +439,15 @@ export const storedUserPreferencesProps: Required<{
     validator: z.boolean().default(false),
     type: 'boolean',
   },
-  enableAIFeatures: {
+  enableGenAIFeatures: {
     ui: true,
     cli: true,
     global: true,
     description: {
       short: 'Enable AI Features',
       long: 'Allow the use of AI features in Compass which make requests to 3rd party services. These features are currently experimental and offered as a preview to only a limited number of users.',
     },
-    deriveValue: deriveNetworkTrafficOptionState('enableAIFeatures'),
+    deriveValue: deriveNetworkTrafficOptionState('enableGenAIFeatures'),
     validator: z.boolean().default(true),
     type: 'boolean',
   },
@@ -1064,7 +1064,7 @@ export class Preferences {
     if (!showedNetworkOptIn) {
       await this.savePreferences({
         autoUpdates: true,
-        enableAIFeatures: true,
+        enableGenAIFeatures: true,
         enableMaps: true,
         trackUsageStatistics: true,
         enableFeedbackPanel: true,

packages/compass-preferences-model/src/utils.ts

@@ -1,5 +1,5 @@
 import preferences, { preferencesAccess, usePreference } from '.';
-import type { ParsedGlobalPreferencesResult } from '.';
+import type { AllPreferences, ParsedGlobalPreferencesResult } from '.';
 import { setupPreferences } from './setup-preferences';
 import { UserStorage } from './storage';
 import type { ReactHooks } from './react';
@@ -41,21 +41,58 @@ export function capMaxTimeMSAtPreferenceLimit<T>(value: T): T | number {
   return value;
 }
 
+/**
+ * Helper method to check whether or not AI feature is enabled in Compass. The
+ * feature is considered enabled if:
+ *  - AI feature flag is enabled
+ *  - config preference that controls AI is enabled
+ *  - either mms backend rollout enabled feature for the compass user or special
+ *    option to bypass the check is passed
+ */
+export function isAIFeatureEnabled(
+  preferences: Pick<
+    AllPreferences,
+    | 'enableGenAIFeatures'
+    | 'enableGenAIExperience'
+    | 'cloudFeatureRolloutAccess'
+    | 'enableAIWithoutRolloutAccess'
+  > = preferencesAccess.getPreferences()
+) {
+  const {
+    // a "kill switch" property from configuration file to be able to disable
+    // feature in global config
+    enableGenAIFeatures,
+    // feature flag
+    enableGenAIExperience,
+    // based on mms backend rollout response
+    cloudFeatureRolloutAccess,
+    // feature flag to bypass rollout access check
+    enableAIWithoutRolloutAccess,
+  } = preferences;
+  return (
+    enableGenAIFeatures &&
+    enableGenAIExperience &&
+    (!!cloudFeatureRolloutAccess?.GEN_AI_COMPASS ||
+      enableAIWithoutRolloutAccess)
+  );
+}
+
 export function useIsAIFeatureEnabled(React: ReactHooks) {
+  const enableGenAIFeatures = usePreference('enableGenAIFeatures', React);
+  const enableGenAIExperience = usePreference('enableGenAIExperience', React);
+  const cloudFeatureRolloutAccess = usePreference(
+    'cloudFeatureRolloutAccess',
+    React
+  );
   const enableAIWithoutRolloutAccess = usePreference(
     'enableAIWithoutRolloutAccess',
     React
   );
-  const enableAIExperience = usePreference('enableAIExperience', React);
-  const isAIFeatureEnabled = usePreference(
-    'cloudFeatureRolloutAccess',
-    React
-  )?.GEN_AI_COMPASS;
-  const enableAIFeatures = usePreference('enableAIFeatures', React);
 
-  return (
-    enableAIExperience &&
-    (enableAIWithoutRolloutAccess || isAIFeatureEnabled) &&
-    enableAIFeatures
-  );
+  return isAIFeatureEnabled({
+    enableGenAIFeatures,
+    enableGenAIExperience,
+    cloudFeatureRolloutAccess,
+    enableAIWithoutRolloutAccess,
+  });
 }

packages/compass-query-bar/src/components/query-bar.spec.tsx

@@ -125,8 +125,8 @@ describe('QueryBar Component', function () {
     beforeEach(function () {
       sandbox = sinon.createSandbox();
       sandbox.stub(preferencesAccess, 'getPreferences').returns({
-        enableAIExperience: true,
-        enableAIFeatures: true,
+        enableGenAIExperience: true,
+        enableGenAIFeatures: true,
         cloudFeatureRolloutAccess: {
           GEN_AI_COMPASS: true,
         },
@@ -172,14 +172,14 @@ describe('QueryBar Component', function () {
     });
   });
 
-  describe('with enableAIExperience ai disabled', function () {
+  describe('with enableGenAIExperience ai disabled', function () {
     let sandbox: sinon.SinonSandbox;
 
     beforeEach(function () {
       sandbox = sinon.createSandbox();
       sandbox.stub(preferencesAccess, 'getPreferences').returns({
-        enableAIExperience: false,
-        enableAIFeatures: true,
+        enableGenAIExperience: false,
+        enableGenAIFeatures: true,
         cloudFeatureRolloutAccess: {
           GEN_AI_COMPASS: true,
         },
@@ -198,14 +198,14 @@ describe('QueryBar Component', function () {
     });
   });
 
-  describe('with enableAIFeatures ai disabled', function () {
+  describe('with enableGenAIFeatures ai disabled', function () {
     let sandbox: sinon.SinonSandbox;
 
     beforeEach(function () {
       sandbox = sinon.createSandbox();
       sandbox.stub(preferencesAccess, 'getPreferences').returns({
-        enableAIExperience: true,
-        enableAIFeatures: false,
+        enableGenAIExperience: true,
+        enableGenAIFeatures: false,
         cloudFeatureRolloutAccess: {
           GEN_AI_COMPASS: true,
         },

packages/compass-settings/src/components/modal.tsx

@@ -19,6 +19,7 @@ import Sidebar from './sidebar';
 import { saveSettings, closeModal } from '../stores/settings';
 import type { RootState } from '../stores';
 import { getUserInfo } from '../stores/atlas-login';
+import { useIsAIFeatureEnabled } from 'compass-preferences-model';
 
 type Settings = {
   name: string;
@@ -27,7 +28,7 @@ type Settings = {
 
 type SettingsModalProps = {
   isOpen: boolean;
-  isOIDCTabEnabled: boolean;
+  isOIDCEnabled: boolean;
   onMount?: () => void;
   onClose: () => void;
   onSave: () => void;
@@ -59,9 +60,10 @@ export const SettingsModal: React.FunctionComponent<SettingsModalProps> = ({
   onMount,
   onClose,
   onSave,
-  isOIDCTabEnabled,
+  isOIDCEnabled,
   hasChangedSettings,
 }) => {
+  const aiFeatureEnabled = useIsAIFeatureEnabled(React);
   const onMountRef = useRef(onMount);
 
   useEffect(() => {
@@ -74,7 +76,11 @@ export const SettingsModal: React.FunctionComponent<SettingsModalProps> = ({
     { name: 'Privacy', component: PrivacySettings },
   ];
 
-  if (isOIDCTabEnabled) {
+  if (
+    isOIDCEnabled ||
+    // because oidc options overlap with atlas login used for ai feature
+    aiFeatureEnabled
+  ) {
     settings.push({
       name: 'OIDC (Preview)',
       component: OIDCSettings,
@@ -133,12 +139,7 @@ export default connect(
     return {
       isOpen:
         state.settings.isModalOpen && state.settings.loadingState === 'ready',
-      isOIDCTabEnabled:
-        !!state.settings.settings.enableOidc ||
-        // because oidc options overlap with atlas login used for ai feature
-        !!state.settings.settings.enableAIWithoutRolloutAccess ||
-        !!state.settings.settings.enableAIExperience ||
-        !!state.settings.settings.enableAIFeatures,
+      isOIDCEnabled: !!state.settings.settings.enableOidc,
       hasChangedSettings: state.settings.updatedFields.length > 0,
     };
   },

packages/compass-settings/src/components/settings/feature-preview.tsx

@@ -3,11 +3,8 @@ import SettingsList from './settings-list';
 import {
   usePreference,
   featureFlags,
-  withPreferences,
+  useIsAIFeatureEnabled,
 } from 'compass-preferences-model';
-import type { UserPreferences } from 'compass-preferences-model';
-import { connect } from 'react-redux';
-import type { RootState } from '../../stores';
 import { ConnectedAtlasLoginSettings } from './atlas-login';
 import { css, spacing } from '@mongodb-js/compass-components';
 
@@ -36,33 +33,24 @@ function useShouldShowPreviewFeatures(): boolean {
 
 export function useShouldShowFeaturePreviewSettings(): boolean {
   // We want show the feature preview settings tab if:
+  // - AI feature flag is enabled
   // - there are feature flags in preview stage
   // - or if:
   //   - we are in a development environment or 'showDevFeatureFlags' is explicitly enabled
   //   - and there are feature flags in 'development' stage.
-  //   - AI feature flag is enabled
-  const enableAIWithoutRolloutAccess = usePreference(
-    'enableAIWithoutRolloutAccess',
-    React
-  );
-  const enableAIExperience = usePreference('enableAIExperience', React);
+  const aiFeatureEnabled = useIsAIFeatureEnabled(React);
   const showDevFeatures = useShouldShowDevFeatures();
   const showPreviewFeatures = useShouldShowPreviewFeatures();
-  return (
-    enableAIWithoutRolloutAccess ||
-    enableAIExperience ||
-    showPreviewFeatures ||
-    showDevFeatures
-  );
+
+  return aiFeatureEnabled || showPreviewFeatures || showDevFeatures;
 }
 
 const atlasSettingsContainerStyles = css({
   marginTop: spacing[3],
 });
 
-export const FeaturePreviewSettings: React.FunctionComponent<{
-  showAtlasLoginSettings?: boolean;
-}> = ({ showAtlasLoginSettings }) => {
+export const FeaturePreviewSettings: React.FunctionComponent = () => {
+  const aiFeatureEnabled = useIsAIFeatureEnabled(React);
   const showPreviewFeatures = useShouldShowPreviewFeatures();
   const showDevFeatures = useShouldShowDevFeatures();
 
@@ -73,7 +61,7 @@ export const FeaturePreviewSettings: React.FunctionComponent<{
         your own risk!
       </div>
 
-      {showAtlasLoginSettings && (
+      {aiFeatureEnabled && (
         <div className={atlasSettingsContainerStyles}>
           <ConnectedAtlasLoginSettings></ConnectedAtlasLoginSettings>
         </div>
@@ -92,30 +80,4 @@ export const FeaturePreviewSettings: React.FunctionComponent<{
   );
 };
 
-export default withPreferences(
-  connect(
-    (
-      state: RootState,
-      ownProps: {
-        enableAIExperience?: boolean;
-        enableAIWithoutRolloutAccess?: boolean;
-        cloudFeatureRolloutAccess?: UserPreferences['cloudFeatureRolloutAccess'];
-      }
-    ) => {
-      return {
-        showAtlasLoginSettings:
-          state.settings.settings.enableAIExperience ||
-          ['authenticated', 'in-progress'].includes(state.atlasLogin.status) ||
-          ownProps.enableAIExperience ||
-          ownProps.enableAIWithoutRolloutAccess ||
-          ownProps.cloudFeatureRolloutAccess?.GEN_AI_COMPASS,
-      };
-    }
-  )(FeaturePreviewSettings),
-  [
-    'enableAIExperience',
-    'cloudFeatureRolloutAccess',
-    'enableAIWithoutRolloutAccess',
-  ],
-  React
-);
+export default FeaturePreviewSettings;