feat(compass-collection): Collection Plugin Experimentation Assignment Integration – CLOUDP-333845

packages/compass-collection/src/index.ts

@@ -6,10 +6,14 @@ import {
   dataServiceLocator,
   type DataServiceLocator,
   type DataService,
+  connectionInfoRefLocator,
 } from '@mongodb-js/compass-connections/provider';
 import { collectionModelLocator } from '@mongodb-js/compass-app-stores/provider';
 import type { WorkspacePlugin } from '@mongodb-js/compass-workspaces';
 import { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider';
+import { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
+import { createLoggerLocator } from '@mongodb-js/compass-logging/provider';
+import { preferencesLocator } from 'compass-preferences-model/provider';
 import {
   CollectionWorkspaceTitle,
   CollectionPluginTitleComponent,
@@ -29,6 +33,10 @@ export const WorkspaceTab: WorkspacePlugin<typeof CollectionWorkspaceTitle> = {
       dataService: dataServiceLocator as DataServiceLocator<keyof DataService>,
       collection: collectionModelLocator,
       workspaces: workspacesServiceLocator,
+      experimentationServices: experimentationServiceLocator,
+      connectionInfoRef: connectionInfoRefLocator,
+      logger: createLoggerLocator('COMPASS-COLLECTION'),
+      preferences: preferencesLocator,
     }
   ),
   content: CollectionTab,

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

@@ -5,6 +5,7 @@ import type AppRegistry from '@mongodb-js/compass-app-registry';
 import type { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider';
 import type { CollectionSubtab } from '@mongodb-js/compass-workspaces';
 import type { DataService } from '@mongodb-js/compass-connections/provider';
+import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry';
 
 function isAction<A extends AnyAction>(
   action: AnyAction,
@@ -20,6 +21,7 @@ type CollectionThunkAction<R, A extends AnyAction = AnyAction> = ThunkAction<
     localAppRegistry: AppRegistry;
     dataService: DataService;
     workspaces: ReturnType<typeof workspacesServiceLocator>;
+    experimentationServices: ReturnType<typeof experimentationServiceLocator>;
   },
   A
 >;

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

@@ -6,6 +6,10 @@ import Sinon from 'sinon';
 import AppRegistry from '@mongodb-js/compass-app-registry';
 import { expect } from 'chai';
 import type { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider';
+import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry';
+import type { connectionInfoRefLocator } from '@mongodb-js/compass-connections/provider';
+import { createNoopLogger } from '@mongodb-js/compass-logging/provider';
+import { ReadOnlyPreferenceAccess } from 'compass-preferences-model/provider';
 
 const defaultMetadata = {
   namespace: 'test.foo',
@@ -32,6 +36,32 @@ const mockCollection = {
   },
 };
 
+const mockAtlasConnectionInfo = {
+  current: {
+    id: 'test-connection',
+    title: 'Test Connection',
+    connectionOptions: {
+      connectionString: 'mongodb://localhost:27017',
+    },
+    atlasMetadata: {
+      clusterName: 'test-cluster',
+      projectId: 'test-project',
+      orgId: 'test-org',
+      clusterUniqueId: 'test-cluster-unique-id',
+      clusterType: 'REPLICASET' as const,
+      clusterState: 'IDLE' as const,
+      metricsId: 'test-metrics-id',
+      metricsType: 'replicaSet' as const,
+      regionalBaseUrl: null,
+      instanceSize: 'M10',
+      supports: {
+        globalWrites: false,
+        rollingIndexes: true,
+      },
+    },
+  },
+};
+
 describe('Collection Tab Content store', function () {
   const sandbox = Sinon.createSandbox();
 
@@ -42,7 +72,19 @@ describe('Collection Tab Content store', function () {
 
   const configureStore = async (
     options: Partial<CollectionTabOptions> = {},
-    workspaces: Partial<ReturnType<typeof workspacesServiceLocator>> = {}
+    workspaces: Partial<ReturnType<typeof workspacesServiceLocator>> = {},
+    experimentationServices: Partial<
+      ReturnType<typeof experimentationServiceLocator>
+    > = {},
+    connectionInfoRef: Partial<
+      ReturnType<typeof connectionInfoRefLocator>
+    > = {},
+    logger = createNoopLogger('COMPASS-COLLECTION-TEST'),
+    preferences = new ReadOnlyPreferenceAccess({
+      enableGenAIFeatures: true,
+      enableGenAIFeaturesAtlasOrg: true,
+      cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
+    })
   ) => {
     ({ store, deactivate } = activatePlugin(
       {
@@ -54,6 +96,10 @@ describe('Collection Tab Content store', function () {
         localAppRegistry,
         collection: mockCollection as any,
         workspaces: workspaces as any,
+        experimentationServices: experimentationServices as any,
+        connectionInfoRef: connectionInfoRef as any,
+        logger,
+        preferences,
       },
       { on() {}, cleanup() {} } as any
     ));
@@ -76,11 +122,128 @@ describe('Collection Tab Content store', function () {
       const store = await configureStore(undefined, {
         openCollectionWorkspaceSubtab,
       });
-      store.dispatch(selectTab('Documents'));
+      store.dispatch(selectTab('Documents') as any);
       expect(openCollectionWorkspaceSubtab).to.have.been.calledWith(
         'workspace-tab-id',
         'Documents'
       );
     });
   });
+
+  describe('experimentation integration', function () {
+    it('should assign experiment when Atlas metadata is available', async function () {
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      await configureStore(
+        undefined,
+        {},
+        { assignExperiment },
+        mockAtlasConnectionInfo
+      );
+
+      await waitFor(() => {
+        expect(assignExperiment).to.have.been.calledOnceWith(
+          'mock-data-generator',
+          {
+            team: 'data-explorer',
+          }
+        );
+      });
+    });
+
+    it('should not assign experiment when Atlas metadata is missing', async function () {
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+      const mockConnectionInfoRef = {
+        current: {
+          id: 'test-connection',
+          title: 'Test Connection',
+          connectionOptions: {
+            connectionString: 'mongodb://localhost:27017',
+          },
+          // No atlasMetadata
+        },
+      };
+
+      await configureStore(
+        undefined,
+        {},
+        { assignExperiment },
+        mockConnectionInfoRef
+      );
+
+      // Wait a bit to ensure assignment would have happened if it was going to
+      await new Promise((resolve) => setTimeout(resolve, 50));
+      expect(assignExperiment).to.not.have.been.called;
+    });
+
+    it('should handle missing experimentationServices gracefully and initialize successfully', async function () {
+      const store = await configureStore(
+        undefined,
+        {},
+        undefined, // No experimentationServices provided
+        mockAtlasConnectionInfo
+      );
+
+      // Store should still be functional despite missing experimentationServices
+      await waitFor(() => {
+        expect(store.getState())
+          .to.have.property('metadata')
+          .deep.eq(defaultMetadata);
+      });
+    });
+
+    it('should not assign experiment when AI features are disabled at the org level', async function () {
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      const mockPreferences = new ReadOnlyPreferenceAccess({
+        enableGenAIFeatures: true,
+        enableGenAIFeaturesAtlasOrg: false, // Disabled at org level
+        cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
+      });
+
+      const store = await configureStore(
+        undefined,
+        {},
+        { assignExperiment },
+        mockAtlasConnectionInfo,
+        undefined,
+        mockPreferences
+      );
+
+      // Wait a bit to ensure assignment would have happened if it was going to
+      await new Promise((resolve) => setTimeout(resolve, 50));
+      expect(assignExperiment).to.not.have.been.called;
+
+      // Store should still be functional
+      await waitFor(() => {
+        expect(store.getState())
+          .to.have.property('metadata')
+          .deep.eq(defaultMetadata);
+      });
+    });
+
+    it('should handle assignment errors gracefully', async function () {
+      const assignExperiment = sandbox.spy(() =>
+        Promise.reject(new Error('Assignment failed'))
+      );
+
+      await configureStore(
+        undefined,
+        {},
+        { assignExperiment },
+        mockAtlasConnectionInfo
+      );
+
+      await waitFor(() => {
+        expect(assignExperiment).to.have.been.calledOnce;
+      });
+
+      // Store should still be functional despite assignment error
+      await waitFor(() => {
+        expect(store.getState())
+          .to.have.property('metadata')
+          .deep.eq(defaultMetadata);
+      });
+    });
+  });
 });

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

@@ -9,6 +9,13 @@ import reducer, {
 import type { Collection } from '@mongodb-js/compass-app-stores/provider';
 import type { ActivateHelpers } from '@mongodb-js/compass-app-registry';
 import type { workspacesServiceLocator } from '@mongodb-js/compass-workspaces/provider';
+import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry';
+import type { connectionInfoRefLocator } from '@mongodb-js/compass-connections/provider';
+import type { Logger } from '@mongodb-js/compass-logging/provider';
+import {
+  isAIFeatureEnabled,
+  type PreferencesAccess,
+} from 'compass-preferences-model/provider';
 
 export type CollectionTabOptions = {
   /**
@@ -31,18 +38,29 @@ export type CollectionTabServices = {
   collection: Collection;
   localAppRegistry: AppRegistry;
   workspaces: ReturnType<typeof workspacesServiceLocator>;
+  experimentationServices: ReturnType<typeof experimentationServiceLocator>;
+  connectionInfoRef: ReturnType<typeof connectionInfoRefLocator>;
+  logger: Logger;
+  preferences: PreferencesAccess;
 };
 
 export function activatePlugin(
   { namespace, editViewName, tabId }: CollectionTabOptions,
   services: CollectionTabServices,
   { on, cleanup }: ActivateHelpers
-) {
+): {
+  store: ReturnType<typeof createStore>;
+  deactivate: () => void;
+} {
   const {
     dataService,
     collection: collectionModel,
     localAppRegistry,
     workspaces,
+    experimentationServices,
+    connectionInfoRef,
+    logger,
+    preferences,
   } = services;
 
   if (!collectionModel) {
@@ -64,6 +82,7 @@ export function activatePlugin(
         dataService,
         workspaces,
         localAppRegistry,
+        experimentationServices,
       })
     )
   );
@@ -86,6 +105,28 @@ export function activatePlugin(
 
   void collectionModel.fetchMetadata({ dataService }).then((metadata) => {
     store.dispatch(collectionMetadataFetched(metadata));
+
+    // Assign experiment for Mock Data Generator:
+    // Only assign when experimentationServices.assignExperiment is initialized,
+    // we're connected to Atlas,
+    // and the org-level setting for AI features is enabled
+    if (
+      experimentationServices?.assignExperiment && // Ensures experimentation services are available
+      connectionInfoRef.current?.atlasMetadata?.clusterName && // Ensures we only assign in Atlas
+      isAIFeatureEnabled(preferences.getPreferences()) // Ensures org-level AI features setting is enabled
+    ) {
+      void experimentationServices
+        .assignExperiment('mock-data-generator', {
+          team: 'Atlas Growth',
+        })
+        .catch((error) => {
+          logger.debug('Mock Data Generator experiment assignment failed', {
+            experiment: 'MOCK_DATA_GENERATOR',
+            namespace: namespace,
+            error: error.message,
+          });
+        });
+    }
   });
 
   return {

packages/compass-telemetry/src/experimentation-provider.tsx

@@ -34,7 +34,7 @@ const initialContext: CompassExperimentationProviderContextValue = {
   },
 };
 
-const ExperimentationContext =
+export const ExperimentationContext =
   createContext<CompassExperimentationProviderContextValue>(initialContext);
 
 // Provider component that accepts MMS experiment utils as props

packages/compass-telemetry/src/index.ts

@@ -7,3 +7,4 @@ export type {
 } from './types';
 
 export { CompassExperimentationProvider } from './experimentation-provider';
+export { experimentationServiceLocator } from './provider';

packages/compass-telemetry/src/provider.tsx

@@ -1,9 +1,11 @@
-import React, { useRef } from 'react';
+import React, { useRef, useContext } from 'react';
 import { createServiceLocator } from '@mongodb-js/compass-app-registry';
 import { createTrack, type TelemetryServiceOptions } from './generic-track';
 import { useLogger } from '@mongodb-js/compass-logging/provider';
 import type { TrackFunction } from './types';
 import { TestName } from './growth-experiments';
+import { ExperimentationContext } from './experimentation-provider';
+import type { types } from '@mongodb-js/mdb-experiment-js';
 
 const noop = () => {
   // noop
@@ -47,6 +49,20 @@ export function useTelemetry(): TrackFunction {
   return track;
 }
 
+// Service locator for experimentation services (non-component access)
+export const experimentationServiceLocator = createServiceLocator(
+  function useExperimentationServices(): {
+    assignExperiment: (
+      experimentName: string,
+      options?: types.AssignOptions<string>
+    ) => Promise<types.AsyncStatus | null>;
+  } {
+    const { assignExperiment } = useContext(ExperimentationContext);
+    return { assignExperiment };
+  },
+  'experimentationServiceLocator'
+);
+
 type FirstArgument<F> = F extends (...args: [infer A, ...any]) => any
   ? A
   : F extends { new (...args: [infer A, ...any]): any }