From bed75ddfc3addc752a24eed9fa55bf977eb8bf1a Mon Sep 17 00:00:00 2001 From: gagik Date: Fri, 25 Jul 2025 12:50:52 +0200 Subject: [PATCH 1/4] chore(compass-preferences-model): use optInGenAIFeatures for Compass and DE --- .../atlas-cloud/collection-ai-query.test.ts | 4 +- .../src/atlas-ai-service.spec.ts | 51 +++++++++++++++++++ .../src/atlas-ai-service.ts | 43 +++++++--------- .../src/store/atlas-optin-reducer.spec.ts | 38 ++++++-------- .../src/store/atlas-optin-reducer.ts | 4 +- .../src/compass-web-preferences-access.ts | 2 +- .../src/preferences-schema.tsx | 9 ++-- packages/compass-web/sandbox/index.tsx | 5 +- .../sandbox/sandbox-atlas-sign-in.tsx | 5 +- packages/compass-web/src/preferences.tsx | 2 +- 10 files changed, 100 insertions(+), 63 deletions(-) diff --git a/packages/compass-e2e-tests/tests/atlas-cloud/collection-ai-query.test.ts b/packages/compass-e2e-tests/tests/atlas-cloud/collection-ai-query.test.ts index efbddc5477d..374c50d9d9e 100644 --- a/packages/compass-e2e-tests/tests/atlas-cloud/collection-ai-query.test.ts +++ b/packages/compass-e2e-tests/tests/atlas-cloud/collection-ai-query.test.ts @@ -49,7 +49,7 @@ describe('Collection ai query', function () { true ); await browser.setFeature('enableGenAIFeaturesAtlasOrg', true); - await browser.setFeature('optInDataExplorerGenAIFeatures', true); + await browser.setFeature('optInGenAIFeatures', true); }); describe('on the documents tab', function () { @@ -170,7 +170,7 @@ describe('Collection ai query', function () { true ); await browser.setFeature('enableGenAIFeaturesAtlasOrg', false); - await browser.setFeature('optInDataExplorerGenAIFeatures', true); + await browser.setFeature('optInGenAIFeatures', true); }); it('should not show the gen ai intro button', async function () { 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 46e190b742b..a00c8bc972e 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts @@ -321,6 +321,57 @@ describe('AtlasAiService', function () { }); }); }); + + describe('optIntoGenAIFeatures', function () { + beforeEach(async function () { + // Reset preferences + await preferences.savePreferences({ + optInGenAIFeatures: false, + }); + }); + + if (apiURLPreset === 'cloud') { + it('should make a POST request to cloud endpoint and save preference when cloud preset', async function () { + const fetchStub = sandbox.stub().resolves(makeResponse({})); + global.fetch = fetchStub; + + await atlasAiService.optIntoGenAIFeatures(); + + // Verify fetch was called with correct parameters + expect(fetchStub).to.have.been.calledOnce; + + expect(fetchStub).to.have.been.calledWith( + '/cloud/settings/optInDataExplorerGenAIFeatures', + { + method: 'POST', + headers: { + 'Content-Type': 'application/x-www-form-urlencoded', + Accept: 'application/json', + }, + body: new URLSearchParams([['value', 'true']]), + } + ); + + // Verify preference was saved + const currentPreferences = preferences.getPreferences(); + expect(currentPreferences.optInGenAIFeatures).to.equal(true); + }); + } else { + it('should not make any fetch request and only save preference when admin-api preset', async function () { + const fetchStub = sandbox.stub().resolves(makeResponse({})); + global.fetch = fetchStub; + + await atlasAiService.optIntoGenAIFeatures(); + + // Verify no fetch was called + expect(fetchStub).to.not.have.been.called; + + // Verify preference was saved + const currentPreferences = preferences.getPreferences(); + expect(currentPreferences.optInGenAIFeatures).to.equal(true); + }); + } + }); }); } }); diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index d504c1fcc9a..9591dc31908 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -9,7 +9,6 @@ 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 { signIntoAtlasWithModalPrompt } from './store/atlas-signin-reducer'; import { getStore } from './store/atlas-ai-store'; import { optIntoGenAIWithModalPrompt } from './store/atlas-optin-reducer'; @@ -277,13 +276,7 @@ export class AtlasAiService { } async ensureAiFeatureAccess({ signal }: { signal?: AbortSignal } = {}) { - // 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(signIntoAtlasWithModalPrompt({ signal })); + return getStore().dispatch(optIntoGenAIWithModalPrompt({ signal })); } private getQueryOrAggregationFromUserInput = async ( @@ -391,23 +384,25 @@ export class AtlasAiService { ); } - // Performs a post request to atlas to set the user opt in preference to true. - async optIntoGenAIFeaturesAtlas() { - await this.atlasService.authenticatedFetch( - this.atlasService.cloudEndpoint( - '/settings/optInDataExplorerGenAIFeatures' - ), - { - method: 'POST', - headers: { - 'Content-Type': 'application/x-www-form-urlencoded', - Accept: 'application/json', - }, - body: new URLSearchParams([['value', 'true']]), - } - ); + async optIntoGenAIFeatures() { + if (this.apiURLPreset === 'cloud') { + // Performs a post request to Atlas to set the user opt in preference to true. + await this.atlasService.authenticatedFetch( + this.atlasService.cloudEndpoint( + 'settings/optInDataExplorerGenAIFeatures' + ), + { + method: 'POST', + headers: { + 'Content-Type': 'application/x-www-form-urlencoded', + Accept: 'application/json', + }, + body: new URLSearchParams([['value', 'true']]), + } + ); + } await this.preferences.savePreferences({ - optInDataExplorerGenAIFeatures: true, + optInGenAIFeatures: true, }); } diff --git a/packages/compass-generative-ai/src/store/atlas-optin-reducer.spec.ts b/packages/compass-generative-ai/src/store/atlas-optin-reducer.spec.ts index 9829e033833..947842f8054 100644 --- a/packages/compass-generative-ai/src/store/atlas-optin-reducer.spec.ts +++ b/packages/compass-generative-ai/src/store/atlas-optin-reducer.spec.ts @@ -20,7 +20,7 @@ describe('atlasOptInReducer', function () { beforeEach(async function () { mockPreferences = await createSandboxFromDefaultPreferences(); await mockPreferences.savePreferences({ - optInDataExplorerGenAIFeatures: false, + optInGenAIFeatures: false, }); }); @@ -31,7 +31,7 @@ describe('atlasOptInReducer', function () { describe('optIn', function () { it('should check state and set state to success if already opted in', async function () { const mockAtlasAiService = { - optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }), + optIntoGenAIFeatures: sandbox.stub().resolves({ sub: '1234' }), }; const store = configureStore({ atlasAuthService: {} as any, @@ -45,8 +45,7 @@ describe('atlasOptInReducer', function () { ); void store.dispatch(atlasAiServiceOptedIn()); await store.dispatch(optIn()); - expect(mockAtlasAiService.optIntoGenAIFeaturesAtlas).not.to.have.been - .called; + expect(mockAtlasAiService.optIntoGenAIFeatures).not.to.have.been.called; expect(store.getState().optIn).to.have.nested.property( 'state', 'optin-success' @@ -55,7 +54,7 @@ describe('atlasOptInReducer', function () { it('should start opt in, and set state to success', async function () { const mockAtlasAiService = { - optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }), + optIntoGenAIFeatures: sandbox.stub().resolves({ sub: '1234' }), }; const store = configureStore({ atlasAuthService: {} as any, @@ -69,8 +68,7 @@ describe('atlasOptInReducer', function () { ); void store.dispatch(optIntoGenAIWithModalPrompt()).catch(() => {}); await store.dispatch(optIn()); - expect(mockAtlasAiService.optIntoGenAIFeaturesAtlas).to.have.been - .calledOnce; + expect(mockAtlasAiService.optIntoGenAIFeatures).to.have.been.calledOnce; expect(store.getState().optIn).to.have.nested.property( 'state', 'optin-success' @@ -81,13 +79,13 @@ describe('atlasOptInReducer', function () { beforeEach(async function () { await mockPreferences.savePreferences({ enableGenAIFeaturesAtlasProject: false, - optInDataExplorerGenAIFeatures: true, + optInGenAIFeatures: true, }); }); it('should start the opt in flow', async function () { const mockAtlasAiService = { - optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }), + optIntoGenAIFeatures: sandbox.stub().resolves({ sub: '1234' }), }; const store = configureStore({ atlasAuthService: {} as any, @@ -101,8 +99,7 @@ describe('atlasOptInReducer', function () { ); void store.dispatch(optIntoGenAIWithModalPrompt()).catch(() => {}); await store.dispatch(optIn()); - expect(mockAtlasAiService.optIntoGenAIFeaturesAtlas).to.have.been - .calledOnce; + expect(mockAtlasAiService.optIntoGenAIFeatures).to.have.been.calledOnce; expect(store.getState().optIn).to.have.nested.property( 'state', 'optin-success' @@ -112,9 +109,7 @@ describe('atlasOptInReducer', function () { it('should fail opt in if opt in failed', async function () { const mockAtlasAiService = { - optIntoGenAIFeaturesAtlas: sandbox - .stub() - .rejects(new Error('Whooops!')), + optIntoGenAIFeatures: sandbox.stub().rejects(new Error('Whooops!')), }; const store = configureStore({ atlasAuthService: {} as any, @@ -127,8 +122,7 @@ describe('atlasOptInReducer', function () { // Avoid unhandled rejections. AttemptStateMap.get(attemptId)?.promise.catch(() => {}); await optInPromise; - expect(mockAtlasAiService.optIntoGenAIFeaturesAtlas).to.have.been - .calledOnce; + expect(mockAtlasAiService.optIntoGenAIFeatures).to.have.been.calledOnce; expect(store.getState().optIn).to.have.nested.property('state', 'error'); }); }); @@ -153,7 +147,7 @@ describe('atlasOptInReducer', function () { it('should cancel opt in if opt in is in progress', async function () { const mockAtlasAiService = { - optIntoGenAIFeaturesAtlas: sandbox + optIntoGenAIFeatures: sandbox .stub() .callsFake(({ signal }: { signal: AbortSignal }) => { return new Promise((resolve, reject) => { @@ -183,10 +177,10 @@ describe('atlasOptInReducer', function () { }); }); - describe('optIntoAtlasWithModalPrompt', function () { + describe('optIntoGenAIWithModalPrompt', function () { it('should resolve when user finishes opt in with prompt flow', async function () { const mockAtlasAiService = { - optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }), + optIntoGenAIFeatures: sandbox.stub().resolves({ sub: '1234' }), }; const store = configureStore({ atlasAuthService: {} as any, @@ -203,7 +197,7 @@ describe('atlasOptInReducer', function () { it('should reject if opt in flow fails', async function () { const mockAtlasAiService = { - optIntoGenAIFeaturesAtlas: sandbox.stub().rejects(new Error('Whoops!')), + optIntoGenAIFeatures: sandbox.stub().rejects(new Error('Whoops!')), }; const store = configureStore({ atlasAuthService: {} as any, @@ -226,7 +220,7 @@ describe('atlasOptInReducer', function () { it('should reject if user dismissed the modal', async function () { const mockAtlasAiService = { - optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }), + optIntoGenAIFeatures: sandbox.stub().resolves({ sub: '1234' }), }; const store = configureStore({ atlasAuthService: {} as any, @@ -249,7 +243,7 @@ describe('atlasOptInReducer', function () { it('should reject if provided signal was aborted', async function () { const mockAtlasAiService = { - optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }), + optIntoGenAIFeatures: sandbox.stub().resolves({ sub: '1234' }), }; const store = configureStore({ atlasAuthService: {} as any, diff --git a/packages/compass-generative-ai/src/store/atlas-optin-reducer.ts b/packages/compass-generative-ai/src/store/atlas-optin-reducer.ts index 1378c26a093..f653f815eb4 100644 --- a/packages/compass-generative-ai/src/store/atlas-optin-reducer.ts +++ b/packages/compass-generative-ai/src/store/atlas-optin-reducer.ts @@ -228,7 +228,7 @@ export const optIntoGenAIWithModalPrompt = ({ const { state } = getState().optIn; if ( (state === 'optin-success' || - preferences.getPreferences().optInDataExplorerGenAIFeatures) && + preferences.getPreferences().optInGenAIFeatures) && preferences.getPreferences().enableGenAIFeaturesAtlasProject ) { return Promise.resolve(); @@ -265,7 +265,7 @@ export const optIn = (): GenAIAtlasOptInThunkAction> => { try { throwIfAborted(signal); - await atlasAiService.optIntoGenAIFeaturesAtlas(); + await atlasAiService.optIntoGenAIFeatures(); dispatch(atlasAiServiceOptedIn()); resolve(); } catch (err) { diff --git a/packages/compass-preferences-model/src/compass-web-preferences-access.ts b/packages/compass-preferences-model/src/compass-web-preferences-access.ts index 259390ea180..05b9095a3db 100644 --- a/packages/compass-preferences-model/src/compass-web-preferences-access.ts +++ b/packages/compass-preferences-model/src/compass-web-preferences-access.ts @@ -7,7 +7,7 @@ import { getActiveUser } from './utils'; const editablePreferences: (keyof UserPreferences)[] = [ // Value can change from false to true during allocation / checking - 'optInDataExplorerGenAIFeatures', + 'optInGenAIFeatures', 'cloudFeatureRolloutAccess', // TODO(COMPASS-9353): Provide a standard for updating Compass preferences in web 'enableIndexesGuidanceExp', diff --git a/packages/compass-preferences-model/src/preferences-schema.tsx b/packages/compass-preferences-model/src/preferences-schema.tsx index 93b5a6540d0..fdeed50452c 100644 --- a/packages/compass-preferences-model/src/preferences-schema.tsx +++ b/packages/compass-preferences-model/src/preferences-schema.tsx @@ -85,7 +85,7 @@ export type UserConfigurablePreferences = PermanentFeatureFlags & | 'web-sandbox-atlas-dev' | 'web-sandbox-atlas-qa' | 'web-sandbox-atlas'; - optInDataExplorerGenAIFeatures: boolean; + optInGenAIFeatures: boolean; // Features that are enabled by default in Compass, but are disabled in Data // Explorer enableExplainPlan: boolean; @@ -810,17 +810,16 @@ export const storedUserPreferencesProps: Required<{ .default('atlas'), type: 'string', }, - optInDataExplorerGenAIFeatures: { + optInGenAIFeatures: { ui: true, cli: false, global: false, description: { - short: 'User Opt-in for Data Explorer Gen AI Features', + short: 'User or Client Opt-in for Gen AI Features', }, - validator: z.boolean().default(true), + validator: z.boolean().default(false), type: 'boolean', }, - enableAtlasSearchIndexes: { ui: true, cli: true, diff --git a/packages/compass-web/sandbox/index.tsx b/packages/compass-web/sandbox/index.tsx index b5bb034c0f3..17f676edf2b 100644 --- a/packages/compass-web/sandbox/index.tsx +++ b/packages/compass-web/sandbox/index.tsx @@ -48,7 +48,7 @@ const App = () => { enableGenAIFeaturesAtlasProject, enableGenAISampleDocumentPassingOnAtlasProject, enableGenAIFeaturesAtlasOrg, - optInDataExplorerGenAIFeatures, + optInGenAIFeatures, } = projectParams ?? {}; const atlasServiceSandboxBackendVariant = @@ -135,8 +135,7 @@ const App = () => { isAtlas && !!enableGenAISampleDocumentPassingOnAtlasProject, enableGenAIFeaturesAtlasOrg: isAtlas && !!enableGenAIFeaturesAtlasOrg, - optInDataExplorerGenAIFeatures: - isAtlas && !!optInDataExplorerGenAIFeatures, + optInGenAIFeatures: isAtlas && !!optInGenAIFeatures, enableDataModeling: true, }} onTrack={sandboxTelemetry.track} diff --git a/packages/compass-web/sandbox/sandbox-atlas-sign-in.tsx b/packages/compass-web/sandbox/sandbox-atlas-sign-in.tsx index fd6c5322273..c1c257881c7 100644 --- a/packages/compass-web/sandbox/sandbox-atlas-sign-in.tsx +++ b/packages/compass-web/sandbox/sandbox-atlas-sign-in.tsx @@ -19,7 +19,7 @@ type ProjectParams = { enableGenAIFeaturesAtlasProject: boolean; enableGenAISampleDocumentPassingOnAtlasProject: boolean; enableGenAIFeaturesAtlasOrg: boolean; - optInDataExplorerGenAIFeatures: boolean; + optInGenAIFeatures: boolean; }; type AtlasLoginReturnValue = @@ -129,8 +129,7 @@ export function useAtlasProxySignIn(): AtlasLoginReturnValue { projectId, csrfToken, csrfTime, - optInDataExplorerGenAIFeatures: - isOptedIntoDataExplorerGenAIFeatures, + optInGenAIFeatures: isOptedIntoDataExplorerGenAIFeatures, enableGenAIFeaturesAtlasOrg: genAIFeaturesEnabled, enableGenAISampleDocumentPassingOnAtlasProject: groupEnabledFeatureFlags.includes( diff --git a/packages/compass-web/src/preferences.tsx b/packages/compass-web/src/preferences.tsx index 13cdb1060cb..7347cdd53b4 100644 --- a/packages/compass-web/src/preferences.tsx +++ b/packages/compass-web/src/preferences.tsx @@ -55,7 +55,7 @@ export function useCompassWebPreferences( enableShell: false, enableCreatingNewConnections: false, enableGlobalWrites: false, - optInDataExplorerGenAIFeatures: false, + optInGenAIFeatures: false, enableConnectInNewWindow: false, ...initialPreferences, }) From 617ff7c0bf598672384baf42b888a5f5f2e8db8e Mon Sep 17 00:00:00 2001 From: gagik Date: Mon, 28 Jul 2025 10:53:11 +0200 Subject: [PATCH 2/4] fix: delete atlas login tests --- .../tests/atlas-login.test.ts | 336 ------------------ 1 file changed, 336 deletions(-) delete mode 100644 packages/compass-e2e-tests/tests/atlas-login.test.ts diff --git a/packages/compass-e2e-tests/tests/atlas-login.test.ts b/packages/compass-e2e-tests/tests/atlas-login.test.ts deleted file mode 100644 index 0983a993503..00000000000 --- a/packages/compass-e2e-tests/tests/atlas-login.test.ts +++ /dev/null @@ -1,336 +0,0 @@ -import type { CompassBrowser } from '../helpers/compass-browser'; -import { - init, - cleanup, - screenshotIfFailed, - Selectors, - skipForWeb, - TEST_COMPASS_WEB, - DEFAULT_CONNECTION_NAME_1, -} from '../helpers/compass'; -import type { Compass } from '../helpers/compass'; -import type { OIDCMockProviderConfig } from '@mongodb-js/oidc-mock-provider'; -import { OIDCMockProvider } from '@mongodb-js/oidc-mock-provider'; -import path from 'path'; -import { expect } from 'chai'; -import { createNumbersCollection } from '../helpers/insert-data'; -import { startMockAtlasServiceServer } from '../helpers/atlas-service'; -import type { Telemetry } from '../helpers/telemetry'; -import { startTelemetryServer } from '../helpers/telemetry'; - -const DEFAULT_TOKEN_PAYLOAD = { - expires_in: 3600, - payload: { - groups: ['testgroup'], - sub: 'testuser', - aud: 'resource-server-audience-value', - }, -}; - -function getTestBrowserShellCommand() { - return `${process.execPath} ${path.resolve( - __dirname, - '..', - 'fixtures', - 'curl.js' - )}`; -} - -describe('Atlas Login', function () { - let compass: Compass; - let browser: CompassBrowser; - let oidcMockProvider: OIDCMockProvider; - let getTokenPayload: OIDCMockProviderConfig['getTokenPayload']; - let stopMockAtlasServer: () => Promise; - let numberOfOIDCAuthRequests = 0; - - before(async function () { - skipForWeb(this, 'atlas-login not supported in compass-web'); - - // Start a mock server to pass an ai response. - const { endpoint, stop } = await startMockAtlasServiceServer(); - stopMockAtlasServer = stop; - process.env.COMPASS_ATLAS_SERVICE_UNAUTH_BASE_URL_OVERRIDE = endpoint; - - function isAuthorised(req: { headers: { authorization?: string } }) { - const [, token] = req.headers.authorization?.split(' ') ?? []; - // We can't check that the issued token is the one received by oidc-plugin - // so we are only checking that it was passed and assuming that it's good - // enough of a validation - return !!token; - } - - oidcMockProvider = await OIDCMockProvider.create({ - getTokenPayload(metadata) { - return getTokenPayload(metadata); - }, - overrideRequestHandler(_url, req, res) { - const url = new URL(_url); - - switch (url.pathname) { - case '/auth-portal-redirect': - res.statusCode = 307; - res.setHeader('Location', url.searchParams.get('fromURI') ?? ''); - res.end(); - break; - case '/authorize': - numberOfOIDCAuthRequests += 1; - break; - case '/v1/userinfo': - if (isAuthorised(req)) { - res.statusCode = 200; - res.write( - JSON.stringify({ - sub: Date.now().toString(32), - firstName: 'First', - lastName: 'Last', - primaryEmail: 'test@example.com', - login: 'test@example.com', - }) - ); - res.end(); - } else { - res.statusCode = 401; - res.end(); - } - break; - case '/v1/introspect': - res.statusCode = 200; - res.write(JSON.stringify({ active: isAuthorised(req) })); - res.end(); - break; - } - }, - }); - - process.env.COMPASS_CLIENT_ID_OVERRIDE = 'testServer'; - process.env.COMPASS_OIDC_ISSUER_OVERRIDE = oidcMockProvider.issuer; - process.env.COMPASS_ATLAS_AUTH_PORTAL_URL_OVERRIDE = `${oidcMockProvider.issuer}/auth-portal-redirect`; - }); - - beforeEach(async function () { - numberOfOIDCAuthRequests = 0; - - getTokenPayload = () => { - return DEFAULT_TOKEN_PAYLOAD; - }; - - compass = await init(this.test?.fullTitle()); - browser = compass.browser; - await browser.setFeature( - 'browserCommandForOIDCAuth', - getTestBrowserShellCommand() - ); - await browser.setupDefaultConnections(); - }); - - afterEach(async function () { - await browser.setFeature('browserCommandForOIDCAuth', undefined); - await screenshotIfFailed(compass, this.currentTest); - await cleanup(compass); - }); - - after(async function () { - if (TEST_COMPASS_WEB) { - return; - } - - await oidcMockProvider?.close(); - delete process.env.COMPASS_CLIENT_ID_OVERRIDE; - delete process.env.COMPASS_OIDC_ISSUER_OVERRIDE; - delete process.env.COMPASS_ATLAS_AUTH_PORTAL_URL_OVERRIDE; - - await stopMockAtlasServer(); - delete process.env.COMPASS_ATLAS_SERVICE_UNAUTH_BASE_URL_OVERRIDE; - }); - - describe('in settings', function () { - it('should sign in user when clicking on "Log in with Atlas" button', async function () { - await browser.openSettingsModal('ai'); - - await browser.clickVisible(Selectors.LogInWithAtlasButton); - - const loginStatus = browser.$(Selectors.AtlasLoginStatus); - await browser.waitUntil(async () => { - return ( - (await loginStatus.getText()).trim() === - 'Logged in with Atlas account test@example.com' - ); - }); - expect(numberOfOIDCAuthRequests).to.eq(1); - }); - - describe('telemetry', () => { - let telemetry: Telemetry; - - before(async function () { - telemetry = await startTelemetryServer(); - }); - - after(async function () { - await telemetry.stop(); - }); - - it('should send identify after the user has logged in', async function () { - const atlasUserIdBefore = await browser.getFeature( - 'telemetryAtlasUserId' - ); - expect(atlasUserIdBefore).to.not.exist; - - await browser.openSettingsModal('ai'); - - await browser.clickVisible(Selectors.LogInWithAtlasButton); - - const loginStatus = browser.$(Selectors.AtlasLoginStatus); - await browser.waitUntil(async () => { - return ( - (await loginStatus.getText()).trim() === - 'Logged in with Atlas account test@example.com' - ); - }); - - const atlasUserIdAfter = await browser.getFeature( - 'telemetryAtlasUserId' - ); - expect(atlasUserIdAfter).to.be.a('string'); - - const identify = telemetry - .events() - .find((entry) => entry.type === 'identify'); - expect(identify.traits.platform).to.equal(process.platform); - expect(identify.traits.arch).to.match(/^(x64|arm64)$/); - }); - }); - - it('should sign out user when "Disconnect" clicked', async function () { - await browser.openSettingsModal('ai'); - await browser.clickVisible(Selectors.LogInWithAtlasButton); - - const loginStatus = browser.$(Selectors.AtlasLoginStatus); - - await browser.waitUntil(async () => { - return ( - (await loginStatus.getText()).trim() === - 'Logged in with Atlas account test@example.com' - ); - }); - - await browser.clickVisible(Selectors.DisconnectAtlasAccountButton); - - await browser.waitUntil(async () => { - return (await loginStatus.getText()).includes( - 'This is a feature powered by generative AI, and may give inaccurate responses' - ); - }); - }); - - it('should sign in user when disconnected and clicking again on "Log in with Atlas" button', async function () { - await browser.openSettingsModal('ai'); - await browser.clickVisible(Selectors.LogInWithAtlasButton); - - let loginStatus = browser.$(Selectors.AtlasLoginStatus); - - await browser.waitUntil(async () => { - return ( - (await loginStatus.getText()).trim() === - 'Logged in with Atlas account test@example.com' - ); - }); - - await browser.clickVisible(Selectors.DisconnectAtlasAccountButton); - - await browser.clickVisible(Selectors.LogInWithAtlasButton); - - loginStatus = browser.$(Selectors.AtlasLoginStatus); - await browser.waitUntil(async () => { - return ( - (await loginStatus.getText()).trim() === - 'Logged in with Atlas account test@example.com' - ); - }); - expect(numberOfOIDCAuthRequests).to.eq(2); - }); - - it('should show toast with error if sign in failed', async function () { - getTokenPayload = () => { - return Promise.reject(new Error('Auth failed')); - }; - - await browser.openSettingsModal('ai'); - await browser.clickVisible(Selectors.LogInWithAtlasButton); - - const errorToast = browser.$(Selectors.AtlasLoginErrorToast); - await errorToast.waitForDisplayed(); - - expect(await errorToast.getText()).to.match( - /Sign in failed\n+unexpected HTTP response status code.+Auth failed/ - ); - }); - }); - - describe('in CRUD view', function () { - beforeEach(async function () { - await createNumbersCollection(); - await browser.disconnectAll(); - await browser.connectToDefaults(); - await browser.navigateToCollectionTab( - DEFAULT_CONNECTION_NAME_1, - 'test', - 'numbers', - 'Documents' - ); - }); - - it('should not show AI input if sign in flow was not finished', async function () { - getTokenPayload = () => { - return new Promise(() => {}); - }; - - const generateQueryButton = browser.$('button*=Generate query'); - await browser.clickVisible(generateQueryButton); - - await browser.clickVisible(Selectors.LogInWithAtlasModalButton); - - // Because leafygreen doesn't render a button there and we don't have any - // control over it - await browser.clickVisible('span=Not now'); - - const aiInput = browser.$(Selectors.GenAITextInput); - expect(await aiInput.isExisting()).to.eq(false); - expect(await generateQueryButton.isDisplayed()).to.eq(true); - }); - }); - - describe('in Aggregation Builder view', function () { - beforeEach(async function () { - await createNumbersCollection(); - await browser.disconnectAll(); - await browser.connectToDefaults(); - await browser.navigateToCollectionTab( - DEFAULT_CONNECTION_NAME_1, - 'test', - 'numbers', - 'Aggregations' - ); - }); - - it('should not show AI input if sign in flow was not finished', async function () { - getTokenPayload = () => { - return new Promise(() => {}); - }; - - const generateQueryButton = browser.$('button*=Generate aggregation'); - await browser.clickVisible(generateQueryButton); - - await browser.clickVisible(Selectors.LogInWithAtlasModalButton); - - // Because leafygreen doesn't render a button there and we don't have any - // control over it - await browser.clickVisible('span=Not now'); - - const aiInput = browser.$(Selectors.GenAITextInput); - expect(await aiInput.isExisting()).to.eq(false); - expect(await generateQueryButton.isDisplayed()).to.eq(true); - }); - }); -}); From 3df8f4a5794fbac91b5de39ce38528598d5d84f8 Mon Sep 17 00:00:00 2001 From: gagik Date: Mon, 28 Jul 2025 11:08:51 +0200 Subject: [PATCH 3/4] fix: cleanup test --- .../src/atlas-ai-service.spec.ts | 36 +++++++------------ 1 file changed, 13 insertions(+), 23 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 a00c8bc972e..be807527b82 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts @@ -330,13 +330,14 @@ describe('AtlasAiService', function () { }); }); - if (apiURLPreset === 'cloud') { - it('should make a POST request to cloud endpoint and save preference when cloud preset', async function () { - const fetchStub = sandbox.stub().resolves(makeResponse({})); - global.fetch = fetchStub; + it('should save preference when cloud preset', async function () { + const fetchStub = sandbox.stub().resolves(makeResponse({})); + global.fetch = fetchStub; - await atlasAiService.optIntoGenAIFeatures(); + await atlasAiService.optIntoGenAIFeatures(); + // In Data Explorer, make a POST request to cloud endpoint and save preference + if (apiURLPreset === 'cloud') { // Verify fetch was called with correct parameters expect(fetchStub).to.have.been.calledOnce; @@ -351,26 +352,15 @@ describe('AtlasAiService', function () { body: new URLSearchParams([['value', 'true']]), } ); - - // Verify preference was saved - const currentPreferences = preferences.getPreferences(); - expect(currentPreferences.optInGenAIFeatures).to.equal(true); - }); - } else { - it('should not make any fetch request and only save preference when admin-api preset', async function () { - const fetchStub = sandbox.stub().resolves(makeResponse({})); - global.fetch = fetchStub; - - await atlasAiService.optIntoGenAIFeatures(); - - // Verify no fetch was called + } else { + // In Compass, no fetch is made, only stored locally expect(fetchStub).to.not.have.been.called; + } - // Verify preference was saved - const currentPreferences = preferences.getPreferences(); - expect(currentPreferences.optInGenAIFeatures).to.equal(true); - }); - } + // Verify preference was saved + const currentPreferences = preferences.getPreferences(); + expect(currentPreferences.optInGenAIFeatures).to.equal(true); + }); }); }); } From 44e5585b7faa066dfbefc9db53bf78be65fb5ea9 Mon Sep 17 00:00:00 2001 From: gagik Date: Mon, 4 Aug 2025 16:56:01 +0200 Subject: [PATCH 4/4] chore: use a feature flag --- .../tests/atlas-login.test.ts | 336 ++++++++++++++++++ .../src/atlas-ai-service.spec.ts | 7 + .../src/atlas-ai-service.ts | 13 +- .../src/feature-flags.ts | 8 + 4 files changed, 363 insertions(+), 1 deletion(-) create mode 100644 packages/compass-e2e-tests/tests/atlas-login.test.ts diff --git a/packages/compass-e2e-tests/tests/atlas-login.test.ts b/packages/compass-e2e-tests/tests/atlas-login.test.ts new file mode 100644 index 00000000000..0983a993503 --- /dev/null +++ b/packages/compass-e2e-tests/tests/atlas-login.test.ts @@ -0,0 +1,336 @@ +import type { CompassBrowser } from '../helpers/compass-browser'; +import { + init, + cleanup, + screenshotIfFailed, + Selectors, + skipForWeb, + TEST_COMPASS_WEB, + DEFAULT_CONNECTION_NAME_1, +} from '../helpers/compass'; +import type { Compass } from '../helpers/compass'; +import type { OIDCMockProviderConfig } from '@mongodb-js/oidc-mock-provider'; +import { OIDCMockProvider } from '@mongodb-js/oidc-mock-provider'; +import path from 'path'; +import { expect } from 'chai'; +import { createNumbersCollection } from '../helpers/insert-data'; +import { startMockAtlasServiceServer } from '../helpers/atlas-service'; +import type { Telemetry } from '../helpers/telemetry'; +import { startTelemetryServer } from '../helpers/telemetry'; + +const DEFAULT_TOKEN_PAYLOAD = { + expires_in: 3600, + payload: { + groups: ['testgroup'], + sub: 'testuser', + aud: 'resource-server-audience-value', + }, +}; + +function getTestBrowserShellCommand() { + return `${process.execPath} ${path.resolve( + __dirname, + '..', + 'fixtures', + 'curl.js' + )}`; +} + +describe('Atlas Login', function () { + let compass: Compass; + let browser: CompassBrowser; + let oidcMockProvider: OIDCMockProvider; + let getTokenPayload: OIDCMockProviderConfig['getTokenPayload']; + let stopMockAtlasServer: () => Promise; + let numberOfOIDCAuthRequests = 0; + + before(async function () { + skipForWeb(this, 'atlas-login not supported in compass-web'); + + // Start a mock server to pass an ai response. + const { endpoint, stop } = await startMockAtlasServiceServer(); + stopMockAtlasServer = stop; + process.env.COMPASS_ATLAS_SERVICE_UNAUTH_BASE_URL_OVERRIDE = endpoint; + + function isAuthorised(req: { headers: { authorization?: string } }) { + const [, token] = req.headers.authorization?.split(' ') ?? []; + // We can't check that the issued token is the one received by oidc-plugin + // so we are only checking that it was passed and assuming that it's good + // enough of a validation + return !!token; + } + + oidcMockProvider = await OIDCMockProvider.create({ + getTokenPayload(metadata) { + return getTokenPayload(metadata); + }, + overrideRequestHandler(_url, req, res) { + const url = new URL(_url); + + switch (url.pathname) { + case '/auth-portal-redirect': + res.statusCode = 307; + res.setHeader('Location', url.searchParams.get('fromURI') ?? ''); + res.end(); + break; + case '/authorize': + numberOfOIDCAuthRequests += 1; + break; + case '/v1/userinfo': + if (isAuthorised(req)) { + res.statusCode = 200; + res.write( + JSON.stringify({ + sub: Date.now().toString(32), + firstName: 'First', + lastName: 'Last', + primaryEmail: 'test@example.com', + login: 'test@example.com', + }) + ); + res.end(); + } else { + res.statusCode = 401; + res.end(); + } + break; + case '/v1/introspect': + res.statusCode = 200; + res.write(JSON.stringify({ active: isAuthorised(req) })); + res.end(); + break; + } + }, + }); + + process.env.COMPASS_CLIENT_ID_OVERRIDE = 'testServer'; + process.env.COMPASS_OIDC_ISSUER_OVERRIDE = oidcMockProvider.issuer; + process.env.COMPASS_ATLAS_AUTH_PORTAL_URL_OVERRIDE = `${oidcMockProvider.issuer}/auth-portal-redirect`; + }); + + beforeEach(async function () { + numberOfOIDCAuthRequests = 0; + + getTokenPayload = () => { + return DEFAULT_TOKEN_PAYLOAD; + }; + + compass = await init(this.test?.fullTitle()); + browser = compass.browser; + await browser.setFeature( + 'browserCommandForOIDCAuth', + getTestBrowserShellCommand() + ); + await browser.setupDefaultConnections(); + }); + + afterEach(async function () { + await browser.setFeature('browserCommandForOIDCAuth', undefined); + await screenshotIfFailed(compass, this.currentTest); + await cleanup(compass); + }); + + after(async function () { + if (TEST_COMPASS_WEB) { + return; + } + + await oidcMockProvider?.close(); + delete process.env.COMPASS_CLIENT_ID_OVERRIDE; + delete process.env.COMPASS_OIDC_ISSUER_OVERRIDE; + delete process.env.COMPASS_ATLAS_AUTH_PORTAL_URL_OVERRIDE; + + await stopMockAtlasServer(); + delete process.env.COMPASS_ATLAS_SERVICE_UNAUTH_BASE_URL_OVERRIDE; + }); + + describe('in settings', function () { + it('should sign in user when clicking on "Log in with Atlas" button', async function () { + await browser.openSettingsModal('ai'); + + await browser.clickVisible(Selectors.LogInWithAtlasButton); + + const loginStatus = browser.$(Selectors.AtlasLoginStatus); + await browser.waitUntil(async () => { + return ( + (await loginStatus.getText()).trim() === + 'Logged in with Atlas account test@example.com' + ); + }); + expect(numberOfOIDCAuthRequests).to.eq(1); + }); + + describe('telemetry', () => { + let telemetry: Telemetry; + + before(async function () { + telemetry = await startTelemetryServer(); + }); + + after(async function () { + await telemetry.stop(); + }); + + it('should send identify after the user has logged in', async function () { + const atlasUserIdBefore = await browser.getFeature( + 'telemetryAtlasUserId' + ); + expect(atlasUserIdBefore).to.not.exist; + + await browser.openSettingsModal('ai'); + + await browser.clickVisible(Selectors.LogInWithAtlasButton); + + const loginStatus = browser.$(Selectors.AtlasLoginStatus); + await browser.waitUntil(async () => { + return ( + (await loginStatus.getText()).trim() === + 'Logged in with Atlas account test@example.com' + ); + }); + + const atlasUserIdAfter = await browser.getFeature( + 'telemetryAtlasUserId' + ); + expect(atlasUserIdAfter).to.be.a('string'); + + const identify = telemetry + .events() + .find((entry) => entry.type === 'identify'); + expect(identify.traits.platform).to.equal(process.platform); + expect(identify.traits.arch).to.match(/^(x64|arm64)$/); + }); + }); + + it('should sign out user when "Disconnect" clicked', async function () { + await browser.openSettingsModal('ai'); + await browser.clickVisible(Selectors.LogInWithAtlasButton); + + const loginStatus = browser.$(Selectors.AtlasLoginStatus); + + await browser.waitUntil(async () => { + return ( + (await loginStatus.getText()).trim() === + 'Logged in with Atlas account test@example.com' + ); + }); + + await browser.clickVisible(Selectors.DisconnectAtlasAccountButton); + + await browser.waitUntil(async () => { + return (await loginStatus.getText()).includes( + 'This is a feature powered by generative AI, and may give inaccurate responses' + ); + }); + }); + + it('should sign in user when disconnected and clicking again on "Log in with Atlas" button', async function () { + await browser.openSettingsModal('ai'); + await browser.clickVisible(Selectors.LogInWithAtlasButton); + + let loginStatus = browser.$(Selectors.AtlasLoginStatus); + + await browser.waitUntil(async () => { + return ( + (await loginStatus.getText()).trim() === + 'Logged in with Atlas account test@example.com' + ); + }); + + await browser.clickVisible(Selectors.DisconnectAtlasAccountButton); + + await browser.clickVisible(Selectors.LogInWithAtlasButton); + + loginStatus = browser.$(Selectors.AtlasLoginStatus); + await browser.waitUntil(async () => { + return ( + (await loginStatus.getText()).trim() === + 'Logged in with Atlas account test@example.com' + ); + }); + expect(numberOfOIDCAuthRequests).to.eq(2); + }); + + it('should show toast with error if sign in failed', async function () { + getTokenPayload = () => { + return Promise.reject(new Error('Auth failed')); + }; + + await browser.openSettingsModal('ai'); + await browser.clickVisible(Selectors.LogInWithAtlasButton); + + const errorToast = browser.$(Selectors.AtlasLoginErrorToast); + await errorToast.waitForDisplayed(); + + expect(await errorToast.getText()).to.match( + /Sign in failed\n+unexpected HTTP response status code.+Auth failed/ + ); + }); + }); + + describe('in CRUD view', function () { + beforeEach(async function () { + await createNumbersCollection(); + await browser.disconnectAll(); + await browser.connectToDefaults(); + await browser.navigateToCollectionTab( + DEFAULT_CONNECTION_NAME_1, + 'test', + 'numbers', + 'Documents' + ); + }); + + it('should not show AI input if sign in flow was not finished', async function () { + getTokenPayload = () => { + return new Promise(() => {}); + }; + + const generateQueryButton = browser.$('button*=Generate query'); + await browser.clickVisible(generateQueryButton); + + await browser.clickVisible(Selectors.LogInWithAtlasModalButton); + + // Because leafygreen doesn't render a button there and we don't have any + // control over it + await browser.clickVisible('span=Not now'); + + const aiInput = browser.$(Selectors.GenAITextInput); + expect(await aiInput.isExisting()).to.eq(false); + expect(await generateQueryButton.isDisplayed()).to.eq(true); + }); + }); + + describe('in Aggregation Builder view', function () { + beforeEach(async function () { + await createNumbersCollection(); + await browser.disconnectAll(); + await browser.connectToDefaults(); + await browser.navigateToCollectionTab( + DEFAULT_CONNECTION_NAME_1, + 'test', + 'numbers', + 'Aggregations' + ); + }); + + it('should not show AI input if sign in flow was not finished', async function () { + getTokenPayload = () => { + return new Promise(() => {}); + }; + + const generateQueryButton = browser.$('button*=Generate aggregation'); + await browser.clickVisible(generateQueryButton); + + await browser.clickVisible(Selectors.LogInWithAtlasModalButton); + + // Because leafygreen doesn't render a button there and we don't have any + // control over it + await browser.clickVisible('span=Not now'); + + const aiInput = browser.$(Selectors.GenAITextInput); + expect(await aiInput.isExisting()).to.eq(false); + expect(await generateQueryButton.isDisplayed()).to.eq(true); + }); + }); +}); 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 be807527b82..193a437acff 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts @@ -327,6 +327,13 @@ describe('AtlasAiService', function () { // Reset preferences await preferences.savePreferences({ optInGenAIFeatures: false, + enableUnauthenticatedGenAI: true, + }); + }); + + afterEach(async function () { + await preferences.savePreferences({ + enableUnauthenticatedGenAI: false, }); }); diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index 9591dc31908..8c5a39b1aa1 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -11,6 +11,7 @@ import type { Logger } from '@mongodb-js/compass-logging'; import { EJSON } from 'bson'; import { getStore } from './store/atlas-ai-store'; import { optIntoGenAIWithModalPrompt } from './store/atlas-optin-reducer'; +import { signIntoAtlasWithModalPrompt } from './store/atlas-signin-reducer'; type GenerativeAiInput = { userInput: string; @@ -276,7 +277,17 @@ export class AtlasAiService { } async ensureAiFeatureAccess({ signal }: { signal?: AbortSignal } = {}) { - return getStore().dispatch(optIntoGenAIWithModalPrompt({ signal })); + if (this.preferences.getPreferences().enableUnauthenticatedGenAI) { + return getStore().dispatch(optIntoGenAIWithModalPrompt({ signal })); + } + + // 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(signIntoAtlasWithModalPrompt({ signal })); } private getQueryOrAggregationFromUserInput = async ( diff --git a/packages/compass-preferences-model/src/feature-flags.ts b/packages/compass-preferences-model/src/feature-flags.ts index ca4d0fe0b2d..b45be8e0e61 100644 --- a/packages/compass-preferences-model/src/feature-flags.ts +++ b/packages/compass-preferences-model/src/feature-flags.ts @@ -28,6 +28,7 @@ export type FeatureFlags = { showIndexesGuidanceVariant: boolean; enableContextMenus: boolean; enableSearchActivationProgramP1: boolean; + enableUnauthenticatedGenAI: boolean; }; export const featureFlags: Required<{ @@ -150,6 +151,13 @@ export const featureFlags: Required<{ }, }, + enableUnauthenticatedGenAI: { + stage: 'development', + description: { + short: 'Enable GenAI for unauthenticated users', + }, + }, + /** * Feature flag for CLOUDP-308952. */