From 4dadfb0a09462a9bf90a1e28adf8dcdd369be144 Mon Sep 17 00:00:00 2001
From: Rhys <Anemy@users.noreply.github.com>
Date: Mon, 9 Dec 2024 13:29:10 -0500
Subject: [PATCH] chore(generative-ai): use telemetryAtlasUserId in feature
 enablement check COMPASS-8571 (#6527)

---
 .../src/atlas-ai-service.spec.ts                   |  9 +++------
 .../compass-generative-ai/src/atlas-ai-service.ts  | 10 +++++++---
 packages/compass-web/sandbox/index.tsx             | 14 +++++++++++---
 packages/compass-web/src/connection-storage.tsx    |  2 +-
 4 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts
index ad904fb4a68..1ecba4a08ca 100644
--- a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts
+++ b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts
@@ -15,11 +15,6 @@ const ATLAS_USER = {
   sub: '123',
 };
 
-const PREFERENCES_USER = {
-  id: '1234',
-  createdAt: new Date(),
-};
-
 const BASE_URL = 'http://example.com';
 
 const mockConnectionInfo: ConnectionInfo = {
@@ -68,7 +63,9 @@ describe('AtlasAiService', function () {
   beforeEach(async function () {
     sandbox = Sinon.createSandbox();
     preferences = await createSandboxFromDefaultPreferences();
-    preferences['getPreferencesUser'] = () => PREFERENCES_USER;
+    await preferences.savePreferences({
+      telemetryAtlasUserId: '1234',
+    });
   });
 
   afterEach(function () {
diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts
index ef89e1602fc..5264a5366fc 100644
--- a/packages/compass-generative-ai/src/atlas-ai-service.ts
+++ b/packages/compass-generative-ai/src/atlas-ai-service.ts
@@ -8,7 +8,7 @@ import { AtlasServiceError } from '@mongodb-js/atlas-service/renderer';
 import type { ConnectionInfo } from '@mongodb-js/compass-connections/provider';
 import type { Document } from 'mongodb';
 import type { Logger } from '@mongodb-js/compass-logging';
-import { EJSON } from 'bson';
+import { EJSON, UUID } from 'bson';
 import { signIntoAtlasWithModalPrompt } from './store/atlas-signin-reducer';
 import { getStore } from './store/atlas-ai-store';
 import { optIntoGenAIWithModalPrompt } from './store/atlas-optin-reducer';
@@ -245,7 +245,8 @@ export class AtlasAiService {
       if (urlId === 'user-access') {
         return this.atlasService.cloudEndpoint(
           aiURLConfig[this.apiURLPreset][urlId](
-            this.preferences.getPreferencesUser().id
+            this.preferences.getPreferences().telemetryAtlasUserId ??
+              new UUID().toString()
           )
         );
       }
@@ -264,7 +265,10 @@ export class AtlasAiService {
     const urlConfig = aiURLConfig[this.apiURLPreset][urlId];
     const urlPath =
       typeof urlConfig === 'function'
-        ? urlConfig(this.preferences.getPreferencesUser().id)
+        ? urlConfig(
+            this.preferences.getPreferences().telemetryAtlasUserId ??
+              new UUID().toString()
+          )
         : urlConfig;
 
     return this.atlasService.adminApiEndpoint(urlPath);
diff --git a/packages/compass-web/sandbox/index.tsx b/packages/compass-web/sandbox/index.tsx
index 9404f893c3c..f7fbd05e9f8 100644
--- a/packages/compass-web/sandbox/index.tsx
+++ b/packages/compass-web/sandbox/index.tsx
@@ -2,7 +2,7 @@ import React, { useLayoutEffect } from 'react';
 import ReactDOM from 'react-dom';
 import { resetGlobalCSS, css, Body } from '@mongodb-js/compass-components';
 import { CompassWeb } from '../src/index';
-import { SandboxConnectionStorageProviver } from '../src/connection-storage';
+import { SandboxConnectionStorageProvider } from '../src/connection-storage';
 import { sandboxLogger } from './sandbox-logger';
 import { sandboxTelemetry } from './sandbox-telemetry';
 import { useAtlasProxySignIn } from './sandbox-atlas-sign-in';
@@ -41,6 +41,9 @@ const App = () => {
       ? 'web-sandbox-atlas-dev'
       : 'web-sandbox-atlas';
 
+  const overrideGenAIEnablement =
+    process.env.COMPASS_WEB_GEN_AI_ENABLEMENT === 'true';
+
   useLayoutEffect(() => {
     getMetaEl('csrf-token').setAttribute('content', csrfToken ?? '');
     getMetaEl('csrf-time').setAttribute('content', csrfTime ?? '');
@@ -53,7 +56,7 @@ const App = () => {
   const isAtlas = status === 'signed-in';
 
   return (
-    <SandboxConnectionStorageProviver
+    <SandboxConnectionStorageProvider
       value={isAtlas ? null : sandboxConnectionStorage}
       extraConnectionOptions={
         isAtlas
@@ -77,13 +80,18 @@ const App = () => {
             enableCreatingNewConnections: !isAtlas,
             enableGlobalWrites: isAtlas,
             enableRollingIndexes: isAtlas,
+            enableGenAIFeaturesAtlasProject: isAtlas && overrideGenAIEnablement,
+            enableGenAISampleDocumentPassingOnAtlasProject:
+              isAtlas && overrideGenAIEnablement,
+            enableGenAIFeaturesAtlasOrg: isAtlas && overrideGenAIEnablement,
+            optInDataExplorerGenAIFeatures: isAtlas && overrideGenAIEnablement,
           }}
           onTrack={sandboxTelemetry.track}
           onDebug={sandboxLogger.debug}
           onLog={sandboxLogger.log}
         ></CompassWeb>
       </Body>
-    </SandboxConnectionStorageProviver>
+    </SandboxConnectionStorageProvider>
   );
 };
 
diff --git a/packages/compass-web/src/connection-storage.tsx b/packages/compass-web/src/connection-storage.tsx
index 111e2648a0d..cf3cdf18ec1 100644
--- a/packages/compass-web/src/connection-storage.tsx
+++ b/packages/compass-web/src/connection-storage.tsx
@@ -317,7 +317,7 @@ const SandboxExtraConnectionOptionsContext = React.createContext<
  * non-Atlas deployment
  * @internal
  */
-export const SandboxConnectionStorageProviver = ({
+export const SandboxConnectionStorageProvider = ({
   value,
   extraConnectionOptions,
   children,