diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index 89cdccfaf8f..ef89e1602fc 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/packages/compass-generative-ai/src/atlas-ai-service.ts @@ -10,7 +10,8 @@ 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-signin-store'; +import { getStore } from './store/atlas-ai-store'; +import { optIntoGenAIWithModalPrompt } from './store/atlas-optin-reducer'; type GenerativeAiInput = { userInput: string; @@ -329,6 +330,10 @@ 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 })); } @@ -437,6 +442,26 @@ 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']]), + } + ); + await this.preferences.savePreferences({ + optInDataExplorerGenAIFeatures: true, + }); + } + private validateAIFeatureEnablementResponse( response: any ): asserts response is AIFeatureEnablement { diff --git a/packages/compass-generative-ai/src/components/atlas-signin/ai-signin-banner-image.tsx b/packages/compass-generative-ai/src/components/ai-image-banner.tsx similarity index 99% rename from packages/compass-generative-ai/src/components/atlas-signin/ai-signin-banner-image.tsx rename to packages/compass-generative-ai/src/components/ai-image-banner.tsx index f515d461df0..cfba44e0f6e 100644 --- a/packages/compass-generative-ai/src/components/atlas-signin/ai-signin-banner-image.tsx +++ b/packages/compass-generative-ai/src/components/ai-image-banner.tsx @@ -7,7 +7,7 @@ const bannerStyles = css({ height: 263, }); -export const AISignInImageBanner = () => { +export const AiImageBanner = () => { return ( + {}} + onOptInClick={() => {}} + > + + ); + expect( + screen.getByRole('heading', { + name: 'Use natural language to generate queries and pipelines', + }) + ).to.exist; + }); + it('should show the cancel button', function () { + render( + + {}} + onOptInClick={() => {}} + > + {' '} + + + ); + const button = screen.getByText('Cancel').closest('button'); + expect(button).to.not.match('disabled'); + }); + + it('should show the opt in button enabled when project AI setting is enabled', async function () { + await mockPreferences.savePreferences({ + enableGenAIFeaturesAtlasProject: true, + }); + render( + + {}} + onOptInClick={() => {}} + > + {' '} + + + ); + const button = screen.getByText('Use Natural Language').closest('button'); + expect(button?.getAttribute('aria-disabled')).to.equal('false'); + }); + + it('should disable the opt in button if project AI setting is disabled ', async function () { + await mockPreferences.savePreferences({ + enableGenAIFeaturesAtlasProject: false, + }); + render( + + {}} + onOptInClick={() => {}} + > + {' '} + + + ); + const button = screen.getByText('Use Natural Language').closest('button'); + expect(button?.getAttribute('aria-disabled')).to.equal('true'); + }); +}); diff --git a/packages/compass-generative-ai/src/components/ai-optin-modal.tsx b/packages/compass-generative-ai/src/components/ai-optin-modal.tsx new file mode 100644 index 00000000000..84a9fd894fa --- /dev/null +++ b/packages/compass-generative-ai/src/components/ai-optin-modal.tsx @@ -0,0 +1,150 @@ +import React from 'react'; +import { connect } from 'react-redux'; +import { + Banner, + Body, + Link, + ConfirmationModal, + SpinLoader, + css, + spacing, + H3, + palette, +} from '@mongodb-js/compass-components'; +import { AiImageBanner } from './ai-image-banner'; +import { closeOptInModal, optIn } from '../store/atlas-optin-reducer'; +import type { RootState } from '../store/atlas-ai-store'; +import { usePreference } from 'compass-preferences-model/provider'; + +const GEN_AI_FAQ_LINK = 'https://www.mongodb.com/docs/generative-ai-faq/'; + +type OptInModalProps = { + isOptInModalVisible: boolean; + isOptInInProgress: boolean; + onOptInModalClose: () => void; + onOptInClick: () => void; + projectId?: string; +}; + +const titleStyles = css({ + marginBottom: spacing[400], + marginTop: spacing[400], + marginLeft: spacing[500], + marginRight: spacing[500], + textAlign: 'center', +}); + +const bodyStyles = css({ + marginBottom: spacing[400], + marginTop: spacing[400], + marginLeft: spacing[300], + marginRight: spacing[300], + display: 'flex', + flexDirection: 'column', + alignItems: 'center', + textAlign: 'center', +}); + +const disclaimerStyles = css({ + color: palette.gray.dark1, + marginTop: spacing[400], + marginLeft: spacing[800], + marginRight: spacing[800], +}); + +const bannerStyles = css({ + padding: spacing[400], + marginTop: spacing[400], + textAlign: 'left', +}); +const getButtonText = (isOptInInProgress: boolean) => { + return ( + <> +  Use Natural Language + {isOptInInProgress && ( + <> +   + + + )} + + ); +}; + +export const AIOptInModal: React.FunctionComponent = ({ + isOptInModalVisible, + isOptInInProgress, + onOptInModalClose, + onOptInClick, + projectId, +}) => { + const isProjectAIEnabled = usePreference('enableGenAIFeaturesAtlasProject'); + const PROJECT_SETTINGS_LINK = projectId + ? window.location.origin + '/v2/' + projectId + '#/settings/groupSettings' + : null; + + const onConfirmClick = () => { + if (isOptInInProgress) { + return; + } + onOptInClick(); + }; + return ( + + + +

+ Use natural language to generate queries and pipelines +

+ Atlas users can now quickly create queries and aggregations with + MongoDB's  intelligent AI-powered feature, available today. + + {isProjectAIEnabled + ? 'AI features are enabled for project users with data access.' + : 'AI features are disabled for project users.'}{' '} + Project Owners can change this setting in the{' '} + {PROJECT_SETTINGS_LINK !== null ? ( + + AI features + + ) : ( + 'AI features ' + )} + section. + +
+ This is a feature powered by generative AI, and may give inaccurate + responses. Please see our{' '} + + FAQ + {' '} + for more information. +
+ + + ); +}; + +export default connect( + (state: RootState) => { + return { + isOptInModalVisible: state.optIn.isModalOpen, + isOptInInProgress: state.optIn.state === 'in-progress', + }; + }, + { onOptInModalClose: closeOptInModal, onOptInClick: optIn } +)(AIOptInModal); diff --git a/packages/compass-generative-ai/src/components/atlas-signin/ai-signin-modal.tsx b/packages/compass-generative-ai/src/components/ai-signin-modal.tsx similarity index 84% rename from packages/compass-generative-ai/src/components/atlas-signin/ai-signin-modal.tsx rename to packages/compass-generative-ai/src/components/ai-signin-modal.tsx index a199945b067..c5ef7a10796 100644 --- a/packages/compass-generative-ai/src/components/atlas-signin/ai-signin-modal.tsx +++ b/packages/compass-generative-ai/src/components/ai-signin-modal.tsx @@ -10,9 +10,9 @@ import { spacing, useDarkMode, } from '@mongodb-js/compass-components'; -import { AISignInImageBanner } from './ai-signin-banner-image'; -import type { AtlasSignInState } from '../../store/atlas-signin-reducer'; -import { closeSignInModal, signIn } from '../../store/atlas-signin-reducer'; +import { AiImageBanner } from './ai-image-banner'; +import { closeSignInModal, signIn } from '../store/atlas-signin-reducer'; +import type { RootState } from '../store/atlas-ai-store'; const GEN_AI_FAQ_LINK = 'https://www.mongodb.com/docs/generative-ai-faq/'; @@ -30,7 +30,7 @@ const titleStyles = css({ alignItems: 'center', }); -const disclaimer = css({ +const disclaimerStyles = css({ padding: `0 ${spacing[900]}px`, }); @@ -46,7 +46,7 @@ const AISignInModal: React.FunctionComponent = ({ +
This is a feature powered by generative AI, and may give inaccurate responses. Please see our{' '} @@ -55,7 +55,7 @@ const AISignInModal: React.FunctionComponent = ({ for more information.
} - graphic={} + graphic={} title={
Use natural language to generate queries and pipelines @@ -100,10 +100,10 @@ const AISignInModal: React.FunctionComponent = ({ }; export default connect( - (state: AtlasSignInState) => { + (state: RootState) => { return { - isSignInModalVisible: state.isModalOpen, - isSignInInProgress: state.state === 'in-progress', + isSignInModalVisible: state.signIn.isModalOpen, + isSignInInProgress: state.signIn.state === 'in-progress', }; }, { onSignInModalClose: closeSignInModal, onSignInClick: signIn } diff --git a/packages/compass-generative-ai/src/components/atlas-signin/index.tsx b/packages/compass-generative-ai/src/components/atlas-signin/index.tsx deleted file mode 100644 index 1f6ac043c6f..00000000000 --- a/packages/compass-generative-ai/src/components/atlas-signin/index.tsx +++ /dev/null @@ -1,11 +0,0 @@ -import React from 'react'; -import AISignInModal from './ai-signin-modal'; -import { ConfirmationModalArea } from '@mongodb-js/compass-components'; - -export const AtlasSignIn = () => { - return ( - - - - ); -}; diff --git a/packages/compass-generative-ai/src/components/index.ts b/packages/compass-generative-ai/src/components/index.ts index cc65c9df4f4..35991d37b8f 100644 --- a/packages/compass-generative-ai/src/components/index.ts +++ b/packages/compass-generative-ai/src/components/index.ts @@ -3,4 +3,4 @@ export { AIExperienceEntry, createAIPlaceholderHTMLPlaceholder, } from './ai-experience-entry'; -export { AtlasSignIn } from './atlas-signin'; +export { AtlasAiPlugin } from './plugin'; diff --git a/packages/compass-generative-ai/src/components/plugin.tsx b/packages/compass-generative-ai/src/components/plugin.tsx new file mode 100644 index 00000000000..d4795d42dc6 --- /dev/null +++ b/packages/compass-generative-ai/src/components/plugin.tsx @@ -0,0 +1,19 @@ +import React from 'react'; +import AISignInModal from './ai-signin-modal'; +import AIOptInModal from './ai-optin-modal'; +import { ConfirmationModalArea } from '@mongodb-js/compass-components'; + +export interface AtlasAiPluginProps { + projectId?: string; +} + +export const AtlasAiPlugin: React.FunctionComponent = ({ + projectId, +}) => { + return ( + + + + + ); +}; diff --git a/packages/compass-generative-ai/src/index.ts b/packages/compass-generative-ai/src/index.ts index a6f55d1bbec..6e1dae845ad 100644 --- a/packages/compass-generative-ai/src/index.ts +++ b/packages/compass-generative-ai/src/index.ts @@ -1,17 +1,20 @@ import { registerHadronPlugin } from 'hadron-app-registry'; import { atlasAuthServiceLocator } from '@mongodb-js/atlas-service/provider'; - -import { activatePlugin } from './store/atlas-signin-store'; -import { AtlasSignIn } from './components'; +import { AtlasAiPlugin } from './components'; +import { atlasAiServiceLocator } from './provider'; +import { preferencesLocator } from 'compass-preferences-model/provider'; +import { activatePlugin } from './store/atlas-ai-store'; export const CompassGenerativeAIPlugin = registerHadronPlugin( { name: 'CompassGenerativeAI', - component: AtlasSignIn, + component: AtlasAiPlugin, activate: activatePlugin, }, { atlasAuthService: atlasAuthServiceLocator, + atlasAiService: atlasAiServiceLocator, + preferences: preferencesLocator, } ); diff --git a/packages/compass-generative-ai/src/store/atlas-signin-store.ts b/packages/compass-generative-ai/src/store/atlas-ai-store.ts similarity index 51% rename from packages/compass-generative-ai/src/store/atlas-signin-store.ts rename to packages/compass-generative-ai/src/store/atlas-ai-store.ts index 389a92549bd..25e7c9cb4be 100644 --- a/packages/compass-generative-ai/src/store/atlas-signin-store.ts +++ b/packages/compass-generative-ai/src/store/atlas-ai-store.ts @@ -1,31 +1,38 @@ -import { createStore, applyMiddleware } from 'redux'; +import { createStore, applyMiddleware, combineReducers } from 'redux'; import thunk from 'redux-thunk'; -import reducer, { - atlasServiceSignedOut, +import signInReducer, { atlasServiceSignedIn, - atlasServiceTokenRefreshFailed, + atlasServiceSignedOut, + atlasServiceSignInTokenRefreshFailed, } from './atlas-signin-reducer'; +import optInReducer from './atlas-optin-reducer'; import type { AtlasAuthService } from '@mongodb-js/atlas-service/provider'; +import type { AtlasAiService } from '../atlas-ai-service'; +import type { PreferencesAccess } from 'compass-preferences-model'; +import type { AtlasAiPluginProps } from '../components/plugin'; import type { ActivateHelpers } from 'hadron-app-registry'; -let store: CompassGenerativeAIServiceStore; +export let store: CompassGenerativeAIServiceStore; + export function getStore() { if (!store) { throw new Error('CompassGenerativeAIPlugin not activated'); } return store; } +const reducer = combineReducers({ + signIn: signInReducer, + optIn: optInReducer, +}); + +export type RootState = ReturnType; -export type CompassGenerativeAIPluginServices = { - atlasAuthService: AtlasAuthService; -}; export function activatePlugin( - _: Record, - services: CompassGenerativeAIPluginServices, + _initialProps: AtlasAiPluginProps, + services: CompassGenerativeAIExtraArgs, { cleanup }: ActivateHelpers ) { store = configureStore(services); - services.atlasAuthService.on('signed-in', () => { void store.dispatch(atlasServiceSignedIn()); }); @@ -35,18 +42,26 @@ export function activatePlugin( }); services.atlasAuthService.on('token-refresh-failed', () => { - void store.dispatch(atlasServiceTokenRefreshFailed()); + void store.dispatch(atlasServiceSignInTokenRefreshFailed()); }); - return { store, deactivate: cleanup }; } +export type CompassGenerativeAIExtraArgs = { + atlasAuthService: AtlasAuthService; + atlasAiService: AtlasAiService; + preferences: PreferencesAccess; +}; export function configureStore({ atlasAuthService, -}: CompassGenerativeAIPluginServices) { + atlasAiService, + preferences, +}: CompassGenerativeAIExtraArgs) { const store = createStore( reducer, - applyMiddleware(thunk.withExtraArgument({ atlasAuthService })) + applyMiddleware( + thunk.withExtraArgument({ atlasAuthService, atlasAiService, preferences }) + ) ); return store; } 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 new file mode 100644 index 00000000000..9198624ad48 --- /dev/null +++ b/packages/compass-generative-ai/src/store/atlas-optin-reducer.spec.ts @@ -0,0 +1,243 @@ +import Sinon from 'sinon'; +import { expect } from 'chai'; +import { configureStore } from './atlas-ai-store'; +import type { PreferencesAccess } from 'compass-preferences-model'; +import { createSandboxFromDefaultPreferences } from 'compass-preferences-model'; +import { + atlasAiServiceOptedIn, + attemptId, + AttemptStateMap, + cancelOptIn, + closeOptInModal, + optIn, + optIntoGenAIWithModalPrompt, +} from './atlas-optin-reducer'; + +describe('atlasOptInReducer', function () { + const sandbox = Sinon.createSandbox(); + let mockPreferences: PreferencesAccess; + + beforeEach(async function () { + mockPreferences = await createSandboxFromDefaultPreferences(); + await mockPreferences.savePreferences({ + optInDataExplorerGenAIFeatures: false, + }); + }); + + afterEach(function () { + sandbox.reset(); + }); + + 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' }), + }; + const store = configureStore({ + atlasAuthService: {} as any, + atlasAiService: mockAtlasAiService as any, + preferences: mockPreferences, + }); + + expect(store.getState().optIn).to.have.nested.property( + 'state', + 'initial' + ); + void store.dispatch(atlasAiServiceOptedIn()); + await store.dispatch(optIn()); + expect(mockAtlasAiService.optIntoGenAIFeaturesAtlas).not.to.have.been + .called; + expect(store.getState().optIn).to.have.nested.property( + 'state', + 'optin-success' + ); + }); + + it('should start opt in, and set state to success', async function () { + const mockAtlasAiService = { + optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }), + }; + const store = configureStore({ + atlasAuthService: {} as any, + atlasAiService: mockAtlasAiService as any, + preferences: mockPreferences, + }); + + expect(store.getState().optIn).to.have.nested.property( + 'state', + 'initial' + ); + void store.dispatch(optIntoGenAIWithModalPrompt()).catch(() => {}); + await store.dispatch(optIn()); + expect(mockAtlasAiService.optIntoGenAIFeaturesAtlas).to.have.been + .calledOnce; + expect(store.getState().optIn).to.have.nested.property( + 'state', + 'optin-success' + ); + }); + + it('should fail opt in if opt in failed', async function () { + const mockAtlasAiService = { + optIntoGenAIFeaturesAtlas: sandbox + .stub() + .rejects(new Error('Whooops!')), + }; + const store = configureStore({ + atlasAuthService: {} as any, + atlasAiService: mockAtlasAiService as any, + preferences: mockPreferences, + }); + + void store.dispatch(optIntoGenAIWithModalPrompt()).catch(() => {}); + const optInPromise = store.dispatch(optIn()); + // Avoid unhandled rejections. + AttemptStateMap.get(attemptId)?.promise.catch(() => {}); + await optInPromise; + expect(mockAtlasAiService.optIntoGenAIFeaturesAtlas).to.have.been + .calledOnce; + expect(store.getState().optIn).to.have.nested.property('state', 'error'); + }); + }); + + describe('cancelOptIn', function () { + it('should do nothing if no opt in is in progress', function () { + const store = configureStore({ + atlasAuthService: {} as any, + atlasAiService: {} as any, + preferences: mockPreferences, + }); + expect(store.getState().optIn).to.have.nested.property( + 'state', + 'initial' + ); + store.dispatch(cancelOptIn()); + expect(store.getState().optIn).to.have.nested.property( + 'state', + 'initial' + ); + }); + + it('should cancel opt in if opt in is in progress', async function () { + const mockAtlasAiService = { + optIntoGenAIFeaturesAtlas: sandbox + .stub() + .callsFake(({ signal }: { signal: AbortSignal }) => { + return new Promise((resolve, reject) => { + signal.addEventListener('abort', () => { + reject(signal.reason); + }); + }); + }), + }; + + const store = configureStore({ + atlasAuthService: {} as any, + atlasAiService: mockAtlasAiService as any, + preferences: mockPreferences, + }); + + void store.dispatch(optIntoGenAIWithModalPrompt()).catch(() => {}); + + await Promise.all([ + store.dispatch(optIn()), + store.dispatch(cancelOptIn()), + ]); + expect(store.getState().optIn).to.have.nested.property( + 'state', + 'canceled' + ); + }); + }); + + describe('optIntoAtlasWithModalPrompt', function () { + it('should resolve when user finishes opt in with prompt flow', async function () { + const mockAtlasAiService = { + optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }), + }; + const store = configureStore({ + atlasAuthService: {} as any, + atlasAiService: mockAtlasAiService as any, + preferences: mockPreferences, + }); + + const optInPromise = store.dispatch(optIntoGenAIWithModalPrompt()); + await store.dispatch(optIn()); + await optInPromise; + + expect(store.getState().optIn).to.have.property('state', 'optin-success'); + }); + + it('should reject if opt in flow fails', async function () { + const mockAtlasAiService = { + optIntoGenAIFeaturesAtlas: sandbox.stub().rejects(new Error('Whoops!')), + }; + const store = configureStore({ + atlasAuthService: {} as any, + atlasAiService: mockAtlasAiService as any, + preferences: mockPreferences, + }); + + const optInPromise = store.dispatch(optIntoGenAIWithModalPrompt()); + await store.dispatch(optIn()); + + try { + await optInPromise; + throw new Error('Expected optInPromise to throw'); + } catch (err) { + expect(err).to.have.property('message', 'Whoops!'); + } + + expect(store.getState().optIn).to.have.property('state', 'error'); + }); + + it('should reject if user dismissed the modal', async function () { + const mockAtlasAiService = { + optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }), + }; + const store = configureStore({ + atlasAuthService: {} as any, + atlasAiService: mockAtlasAiService as any, + preferences: mockPreferences, + }); + + const optInPromise = store.dispatch(optIntoGenAIWithModalPrompt()); + store.dispatch(closeOptInModal(new Error('This operation was aborted'))); + + try { + await optInPromise; + throw new Error('Expected optInPromise to throw'); + } catch (err) { + expect(err).to.have.property('message', 'This operation was aborted'); + } + + expect(store.getState().optIn).to.have.property('state', 'canceled'); + }); + + it('should reject if provided signal was aborted', async function () { + const mockAtlasAiService = { + optIntoGenAIFeaturesAtlas: sandbox.stub().resolves({ sub: '1234' }), + }; + const store = configureStore({ + atlasAuthService: {} as any, + atlasAiService: mockAtlasAiService as any, + preferences: mockPreferences, + }); + + const c = new AbortController(); + const optInPromise = store.dispatch( + optIntoGenAIWithModalPrompt({ signal: c.signal }) + ); + c.abort(new Error('Aborted from outside')); + + try { + await optInPromise; + throw new Error('Expected optInPromise to throw'); + } catch (err) { + expect(err).to.have.property('message', 'Aborted from outside'); + } + + expect(store.getState().optIn).to.have.property('state', 'canceled'); + }); + }); +}); diff --git a/packages/compass-generative-ai/src/store/atlas-optin-reducer.ts b/packages/compass-generative-ai/src/store/atlas-optin-reducer.ts new file mode 100644 index 00000000000..fbda8cae056 --- /dev/null +++ b/packages/compass-generative-ai/src/store/atlas-optin-reducer.ts @@ -0,0 +1,314 @@ +import type { Action, AnyAction, Reducer } from 'redux'; +import type { ThunkAction } from 'redux-thunk'; +import { throwIfAborted } from '@mongodb-js/compass-utils'; +import type { AtlasAiService } from '../atlas-ai-service'; +import type { PreferencesAccess } from 'compass-preferences-model'; +import type { RootState } from './atlas-ai-store'; +import { isAction } from '../utils/util'; + +type AttemptState = { + id: number; + controller: AbortController; + promise: Promise; + resolve: () => void; + reject: (reason?: any) => void; +}; + +export type AtlasOptInState = { + error: string | null; + isModalOpen: boolean; + attemptId: number | null; + state: 'initial' | 'in-progress' | 'error' | 'canceled' | 'optin-success'; +}; + +export type GenAIAtlasOptInThunkAction< + R, + A extends AnyAction = AnyAction +> = ThunkAction< + R, + RootState, + { atlasAiService: AtlasAiService; preferences: PreferencesAccess }, + A +>; + +export const enum AtlasOptInActions { + OpenOptInModal = 'compass-generative-ai/atlas-optin/OpenOptInModal', + CloseOptInModal = 'compass-generative-ai/atlas-optin/CloseOptInModal', + AttemptStart = 'compass-generative-ai/atlas-optin/AttemptStart', + AttemptEnd = 'compass-generative-ai/atlas-optin/AttemptEnd', + Start = 'compass-generative-ai/atlas-optin/AtlasOptInStart', + OptInSuccess = 'compass-generative-ai/atlas-optin/AtlasOptInSuccess', + Error = 'compass-generative-ai/atlas-optin/AtlasOptInError', + Cancel = 'compass-generative-ai/atlas-optin/AtlasOptInCancel', +} + +export type AtlasOptInOpenModalAction = { + type: AtlasOptInActions.OpenOptInModal; +}; + +export type AtlasOptInCloseModalAction = { + type: AtlasOptInActions.CloseOptInModal; +}; + +export type AtlasOptInAttemptStartAction = { + type: AtlasOptInActions.AttemptStart; + attemptId: number; +}; + +export type AtlasOptInAttemptEndAction = { + type: AtlasOptInActions.AttemptEnd; + attemptId: number; +}; + +export type AtlasOptInStartAction = { + type: AtlasOptInActions.Start; +}; + +export type AtlasOptInSuccessAction = { + type: AtlasOptInActions.OptInSuccess; +}; + +export type AtlasOptInErrorAction = { + type: AtlasOptInActions.Error; + error: string; +}; + +export type AtlasOptInCancelAction = { type: AtlasOptInActions.Cancel }; + +const INITIAL_STATE = { + state: 'initial' as const, + error: null, + isModalOpen: false, + attemptId: null, +}; + +// Exported for testing purposes only. +export const AttemptStateMap = new Map(); + +export let attemptId = 0; + +export function getAttempt(id?: number | null): AttemptState { + if (!id) { + id = ++attemptId; + const controller = new AbortController(); + let resolve; + let reject; + const promise = new Promise((res, rej) => { + resolve = res; + reject = rej; + }); + if (resolve && reject) { + AttemptStateMap.set(id, { + id, + controller, + promise, + resolve: resolve, + reject: reject, + }); + } + } + const attemptState = AttemptStateMap.get(id); + if (!attemptState) { + throw new Error( + 'Trying to get the state for a non-existing opt in attempt' + ); + } + return attemptState; +} + +const optInReducer: Reducer = ( + state = { ...INITIAL_STATE }, + action +) => { + if ( + isAction( + action, + AtlasOptInActions.AttemptStart + ) + ) { + return { + ...state, + attemptId: action.attemptId, + }; + } + + if ( + isAction(action, AtlasOptInActions.AttemptEnd) + ) { + return { + ...state, + attemptId: null, + }; + } + + if (isAction(action, AtlasOptInActions.Start)) { + return { ...state, state: 'in-progress' }; + } + + if ( + isAction(action, AtlasOptInActions.OptInSuccess) + ) { + return { + ...state, + state: 'optin-success', + error: null, + isModalOpen: false, + }; + } + + if (isAction(action, AtlasOptInActions.Error)) { + return { + ...state, + state: 'error', + error: action.error, + isModalOpen: false, + }; + } + + if (isAction(action, AtlasOptInActions.Cancel)) { + return { ...INITIAL_STATE, state: 'canceled' }; + } + + if ( + isAction( + action, + AtlasOptInActions.OpenOptInModal + ) + ) { + return { ...state, isModalOpen: true }; + } + + if ( + isAction( + action, + AtlasOptInActions.CloseOptInModal + ) + ) { + return { ...state, isModalOpen: false }; + } + + return state; +}; + +const startAttempt = ( + fn: () => void +): GenAIAtlasOptInThunkAction => { + return (dispatch, getState) => { + if (getState().optIn.attemptId) { + throw new Error( + "Can't start opt in with prompt while another opt in attempt is in progress" + ); + } + //if pref set to false then call an opt in function then show it + const attempt = getAttempt(); + dispatch({ type: AtlasOptInActions.AttemptStart, attemptId: attempt.id }); + attempt.promise + .finally(() => { + dispatch({ + type: AtlasOptInActions.AttemptEnd, + attemptId: attempt.id, + }); + }) + .catch(() => { + // noop for the promise created by `finally`, original promise rejection + // should be handled by the service user + }); + setTimeout(fn); + return attempt; + }; +}; + +export const optIntoGenAIWithModalPrompt = ({ + signal, +}: { signal?: AbortSignal } = {}): GenAIAtlasOptInThunkAction< + Promise +> => { + return (dispatch, getState, { preferences }) => { + // Nothing to do if we already opted in. + const { state } = getState().optIn; + if ( + state === 'optin-success' || + preferences.getPreferences().optInDataExplorerGenAIFeatures + ) { + return Promise.resolve(); + } + const attempt = dispatch( + startAttempt(() => { + dispatch(openOptInModal()); + }) + ); + signal?.addEventListener('abort', () => { + dispatch(closeOptInModal(signal.reason)); + }); + return attempt.promise; + }; +}; + +export const optIn = (): GenAIAtlasOptInThunkAction> => { + return async (dispatch, getState, { atlasAiService }) => { + if (['in-progress', 'optin-success'].includes(getState().optIn.state)) { + return; + } + const { attemptId } = getState().optIn; + if (attemptId === null) { + return; + } + const { + controller: { signal }, + resolve, + reject, + } = getAttempt(getState().optIn.attemptId); + dispatch({ + type: AtlasOptInActions.Start, + }); + + try { + throwIfAborted(signal); + await atlasAiService.optIntoGenAIFeaturesAtlas(); + dispatch(atlasAiServiceOptedIn()); + resolve(); + } catch (err) { + if (signal.aborted) { + return; + } + dispatch({ + type: AtlasOptInActions.Error, + error: (err as Error).message, + }); + reject(err); + } + }; +}; + +export const openOptInModal = () => { + return { type: AtlasOptInActions.OpenOptInModal }; +}; + +export const closeOptInModal = ( + reason?: any +): GenAIAtlasOptInThunkAction => { + return (dispatch) => { + dispatch(cancelOptIn(reason)); + dispatch({ type: AtlasOptInActions.CloseOptInModal }); + }; +}; + +export const cancelOptIn = (reason?: any): GenAIAtlasOptInThunkAction => { + return (dispatch, getState) => { + // Can't cancel opt in after the flow was finished indicated by current + // attempt id being set to null. + if (getState().optIn.attemptId === null) { + return; + } + const attempt = getAttempt(getState().optIn.attemptId); + attempt.controller.abort(); + attempt.reject(reason ?? attempt.controller.signal.reason); + dispatch({ type: AtlasOptInActions.Cancel }); + }; +}; + +export const atlasAiServiceOptedIn = () => ({ + type: AtlasOptInActions.OptInSuccess, +}); + +export default optInReducer; diff --git a/packages/compass-generative-ai/src/store/atlas-signin-reducer.spec.ts b/packages/compass-generative-ai/src/store/atlas-signin-reducer.spec.ts index 3a036cb536c..e8b0a5f84a2 100644 --- a/packages/compass-generative-ai/src/store/atlas-signin-reducer.spec.ts +++ b/packages/compass-generative-ai/src/store/atlas-signin-reducer.spec.ts @@ -10,10 +10,17 @@ import { closeSignInModal, atlasServiceSignedIn, } from './atlas-signin-reducer'; -import { configureStore } from './atlas-signin-store'; +import { configureStore } from './atlas-ai-store'; +import type { PreferencesAccess } from 'compass-preferences-model'; +import { createSandboxFromDefaultPreferences } from 'compass-preferences-model'; describe('atlasSignInReducer', function () { const sandbox = Sinon.createSandbox(); + let mockPreferences: PreferencesAccess; + + beforeEach(async function () { + mockPreferences = await createSandboxFromDefaultPreferences(); + }); afterEach(function () { sandbox.reset(); @@ -26,13 +33,21 @@ describe('atlasSignInReducer', function () { }; const store = configureStore({ atlasAuthService: mockAtlasService as any, + atlasAiService: mockAtlasService as any, + preferences: mockPreferences, }); - expect(store.getState()).to.have.nested.property('state', 'initial'); + expect(store.getState().signIn).to.have.nested.property( + 'state', + 'initial' + ); void store.dispatch(atlasServiceSignedIn()); await store.dispatch(signIn()); expect(mockAtlasService.signIn).not.to.have.been.called; - expect(store.getState()).to.have.nested.property('state', 'success'); + expect(store.getState().signIn).to.have.nested.property( + 'state', + 'success' + ); }); it('should start sign in, and set state to success', async function () { @@ -41,13 +56,21 @@ describe('atlasSignInReducer', function () { }; const store = configureStore({ atlasAuthService: mockAtlasService as any, + atlasAiService: mockAtlasService as any, + preferences: mockPreferences, }); - expect(store.getState()).to.have.nested.property('state', 'initial'); + expect(store.getState().signIn).to.have.nested.property( + 'state', + 'initial' + ); void store.dispatch(signIntoAtlasWithModalPrompt()).catch(() => {}); await store.dispatch(signIn()); expect(mockAtlasService.signIn).to.have.been.calledOnce; - expect(store.getState()).to.have.nested.property('state', 'success'); + expect(store.getState().signIn).to.have.nested.property( + 'state', + 'success' + ); }); it('should fail sign in if sign in failed', async function () { @@ -56,15 +79,16 @@ describe('atlasSignInReducer', function () { }; const store = configureStore({ atlasAuthService: mockAtlasService as any, + atlasAiService: mockAtlasService as any, + preferences: mockPreferences, }); - void store.dispatch(signIntoAtlasWithModalPrompt()).catch(() => {}); const signInPromise = store.dispatch(signIn()); // Avoid unhandled rejections. AttemptStateMap.get(attemptId)?.promise.catch(() => {}); await signInPromise; expect(mockAtlasService.signIn).to.have.been.calledOnce; - expect(store.getState()).to.have.nested.property('state', 'error'); + expect(store.getState().signIn).to.have.nested.property('state', 'error'); }); }); @@ -72,10 +96,18 @@ describe('atlasSignInReducer', function () { it('should do nothing if no sign in is in progress', function () { const store = configureStore({ atlasAuthService: {} as any, + atlasAiService: {} as any, + preferences: mockPreferences, }); - expect(store.getState()).to.have.nested.property('state', 'initial'); + expect(store.getState().signIn).to.have.nested.property( + 'state', + 'initial' + ); store.dispatch(cancelSignIn()); - expect(store.getState()).to.have.nested.property('state', 'initial'); + expect(store.getState().signIn).to.have.nested.property( + 'state', + 'initial' + ); }); it('should cancel sign in if sign in is in progress', async function () { @@ -92,6 +124,8 @@ describe('atlasSignInReducer', function () { }; const store = configureStore({ atlasAuthService: mockAtlasService as any, + atlasAiService: mockAtlasService as any, + preferences: mockPreferences, }); void store.dispatch(signIntoAtlasWithModalPrompt()).catch(() => {}); @@ -100,7 +134,10 @@ describe('atlasSignInReducer', function () { store.dispatch(signIn()), store.dispatch(cancelSignIn()), ]); - expect(store.getState()).to.have.nested.property('state', 'canceled'); + expect(store.getState().signIn).to.have.nested.property( + 'state', + 'canceled' + ); }); }); @@ -111,13 +148,15 @@ describe('atlasSignInReducer', function () { }; const store = configureStore({ atlasAuthService: mockAtlasService as any, + atlasAiService: mockAtlasService as any, + preferences: mockPreferences, }); const signInPromise = store.dispatch(signIntoAtlasWithModalPrompt()); await store.dispatch(signIn()); await signInPromise; - expect(store.getState()).to.have.property('state', 'success'); + expect(store.getState().signIn).to.have.property('state', 'success'); }); it('should reject if sign in flow fails', async function () { @@ -126,6 +165,8 @@ describe('atlasSignInReducer', function () { }; const store = configureStore({ atlasAuthService: mockAtlasService as any, + atlasAiService: mockAtlasService as any, + preferences: mockPreferences, }); const signInPromise = store.dispatch(signIntoAtlasWithModalPrompt()); @@ -138,7 +179,7 @@ describe('atlasSignInReducer', function () { expect(err).to.have.property('message', 'Whoops!'); } - expect(store.getState()).to.have.property('state', 'error'); + expect(store.getState().signIn).to.have.property('state', 'error'); }); it('should reject if user dismissed the modal', async function () { @@ -147,6 +188,8 @@ describe('atlasSignInReducer', function () { }; const store = configureStore({ atlasAuthService: mockAtlasService as any, + atlasAiService: mockAtlasService as any, + preferences: mockPreferences, }); const signInPromise = store.dispatch(signIntoAtlasWithModalPrompt()); @@ -159,7 +202,7 @@ describe('atlasSignInReducer', function () { expect(err).to.have.property('message', 'This operation was aborted'); } - expect(store.getState()).to.have.property('state', 'canceled'); + expect(store.getState().signIn).to.have.property('state', 'canceled'); }); it('should reject if provided signal was aborted', async function () { @@ -168,6 +211,8 @@ describe('atlasSignInReducer', function () { }; const store = configureStore({ atlasAuthService: mockAtlasService as any, + atlasAiService: mockAtlasService as any, + preferences: mockPreferences, }); const c = new AbortController(); @@ -182,8 +227,7 @@ describe('atlasSignInReducer', function () { } catch (err) { expect(err).to.have.property('message', 'Aborted from outside'); } - - expect(store.getState()).to.have.property('state', 'canceled'); + expect(store.getState().signIn).to.have.property('state', 'canceled'); }); }); }); diff --git a/packages/compass-generative-ai/src/store/atlas-signin-reducer.ts b/packages/compass-generative-ai/src/store/atlas-signin-reducer.ts index 52b15b96920..38a1a9a63be 100644 --- a/packages/compass-generative-ai/src/store/atlas-signin-reducer.ts +++ b/packages/compass-generative-ai/src/store/atlas-signin-reducer.ts @@ -2,13 +2,8 @@ import type { Action, AnyAction, Reducer } from 'redux'; import type { ThunkAction } from 'redux-thunk'; import type { AtlasAuthService } from '@mongodb-js/atlas-service/provider'; import { throwIfAborted } from '@mongodb-js/compass-utils'; - -function isAction( - action: AnyAction, - type: A['type'] -): action is A { - return action.type === type; -} +import type { RootState } from './atlas-ai-store'; +import { isAction } from '../utils/util'; type AttemptState = { id: number; @@ -32,7 +27,7 @@ export type AtlasSignInState = { export type GenAIAtlasSignInThunkAction< R, A extends AnyAction = AnyAction -> = ThunkAction; +> = ThunkAction; export const enum AtlasSignInActions { OpenSignInModal = 'compass-generative-ai/atlas-signin/OpenSignInModal', @@ -43,7 +38,7 @@ export const enum AtlasSignInActions { Success = 'compass-generative-ai/atlas-signin/AtlasSignInSuccess', Error = 'compass-generative-ai/atlas-signin/AtlasSignInError', Cancel = 'compass-generative-ai/atlas-signin/AtlasSignInCancel', - TokenRefreshFailed = 'compass-generative-ai/atlas-signin/TokenRefreshFailed', + SignInTokenRefreshFailed = 'compass-generative-ai/atlas-signin/SignInTokenRefreshFailed', SignedOut = 'compass-generative-ai/atlas-signin/SignedOut', } @@ -79,7 +74,7 @@ export type AtlasSignInErrorAction = { }; export type AtlasSignInTokenRefreshFailedAction = { - type: AtlasSignInActions.TokenRefreshFailed; + type: AtlasSignInActions.SignInTokenRefreshFailed; }; export type AtlasSignInSignedOutAction = { @@ -129,7 +124,7 @@ export function getAttempt(id?: number | null): AttemptState { return attemptState; } -const reducer: Reducer = ( +const signInReducer: Reducer = ( state = { ...INITIAL_STATE }, action ) => { @@ -201,7 +196,7 @@ const reducer: Reducer = ( if ( isAction( action, - AtlasSignInActions.TokenRefreshFailed + AtlasSignInActions.SignInTokenRefreshFailed ) ) { // Only reset state on refresh failed when we are currently successfully @@ -227,7 +222,7 @@ const startAttempt = ( fn: () => void ): GenAIAtlasSignInThunkAction => { return (dispatch, getState) => { - if (getState().attemptId) { + if (getState().signIn.attemptId) { throw new Error( "Can't start sign in with prompt while another sign in attempt is in progress" ); @@ -257,7 +252,7 @@ export const signIntoAtlasWithModalPrompt = ({ > => { return (dispatch, getState) => { // Nothing to do if we already signed in. - const { state } = getState(); + const { state } = getState().signIn; if (state === 'success') { return Promise.resolve(); } @@ -275,10 +270,10 @@ export const signIntoAtlasWithModalPrompt = ({ export const signIn = (): GenAIAtlasSignInThunkAction> => { return async (dispatch, getState, { atlasAuthService }) => { - if (['in-progress', 'authenticated'].includes(getState().state)) { + if (['in-progress', 'authenticated'].includes(getState().signIn.state)) { return; } - const { attemptId } = getState(); + const { attemptId } = getState().signIn; if (attemptId === null) { return; } @@ -286,7 +281,7 @@ export const signIn = (): GenAIAtlasSignInThunkAction> => { controller: { signal }, resolve, reject, - } = getAttempt(getState().attemptId); + } = getAttempt(getState().signIn.attemptId); dispatch({ type: AtlasSignInActions.Start, }); @@ -331,19 +326,19 @@ export const cancelSignIn = ( ): GenAIAtlasSignInThunkAction => { return (dispatch, getState) => { // Can't cancel sign in after the flow was finished indicated by current - // attempt id being set to null. - if (getState().attemptId === null) { + // attempt id being set to null + if (getState().signIn.attemptId === null) { return; } - const attempt = getAttempt(getState().attemptId); + const attempt = getAttempt(getState().signIn.attemptId); attempt.controller.abort(); attempt.reject(reason ?? attempt.controller.signal.reason); dispatch({ type: AtlasSignInActions.Cancel }); }; }; -export const atlasServiceTokenRefreshFailed = () => ({ - type: AtlasSignInActions.TokenRefreshFailed, +export const atlasServiceSignInTokenRefreshFailed = () => ({ + type: AtlasSignInActions.SignInTokenRefreshFailed, }); export const atlasServiceSignedOut = () => ({ @@ -354,4 +349,4 @@ export const atlasServiceSignedIn = () => ({ type: AtlasSignInActions.Success, }); -export default reducer; +export default signInReducer; diff --git a/packages/compass-generative-ai/src/utils/util.ts b/packages/compass-generative-ai/src/utils/util.ts new file mode 100644 index 00000000000..75a77d232eb --- /dev/null +++ b/packages/compass-generative-ai/src/utils/util.ts @@ -0,0 +1,8 @@ +import type { AnyAction } from 'redux'; + +export function isAction( + action: AnyAction, + type: A['type'] +): action is A { + return action.type === type; +} diff --git a/packages/compass-preferences-model/src/preferences-schema.ts b/packages/compass-preferences-model/src/preferences-schema.ts index 53d49eb2052..f3cfdb565a1 100644 --- a/packages/compass-preferences-model/src/preferences-schema.ts +++ b/packages/compass-preferences-model/src/preferences-schema.ts @@ -115,6 +115,15 @@ export type NonUserPreferences = { password?: string; }; +export type AtlasProjectPreferences = { + enableGenAIFeaturesAtlasProject: boolean; + enableGenAISampleDocumentPassingOnAtlasProject: boolean; +}; + +export type AtlasOrgPreferences = { + enableGenAIFeaturesAtlasOrg: boolean; +}; + export type AllPreferences = UserPreferences & CliOnlyPreferences & NonUserPreferences & @@ -213,15 +222,6 @@ export type StoredPreferencesValidator = ReturnType< export type StoredPreferences = z.output; -export type AtlasProjectPreferences = { - enableGenAIFeaturesAtlasProject: boolean; - enableGenAISampleDocumentPassingOnAtlasProject: boolean; -}; - -export type AtlasOrgPreferences = { - enableGenAIFeaturesAtlasOrg: boolean; -}; - // Preference definitions const featureFlagsProps: Required<{ [K in keyof FeatureFlags]: PreferenceDefinition; diff --git a/packages/compass-web/src/entrypoint.tsx b/packages/compass-web/src/entrypoint.tsx index bc5b7736656..b7c679cda4c 100644 --- a/packages/compass-web/src/entrypoint.tsx +++ b/packages/compass-web/src/entrypoint.tsx @@ -30,6 +30,7 @@ import { CompassSchemaPlugin } from '@mongodb-js/compass-schema'; import { CompassIndexesPlugin } from '@mongodb-js/compass-indexes'; import { CompassSchemaValidationPlugin } from '@mongodb-js/compass-schema-validation'; import { CompassGlobalWritesPlugin } from '@mongodb-js/compass-global-writes'; +import { CompassGenerativeAIPlugin } from '@mongodb-js/compass-generative-ai'; import ExplainPlanCollectionTabModal from '@mongodb-js/compass-explain-plan'; import ExportToLanguageCollectionTabModal from '@mongodb-js/compass-export-to-language'; import { @@ -268,7 +269,10 @@ const CompassWeb = ({ enableAggregationBuilderRunPipeline: true, enableAggregationBuilderExtraOptions: true, enableImportExport: false, - enableGenAIFeatures: false, + enableGenAIFeatures: true, + enableGenAIFeaturesAtlasProject: false, + enableGenAISampleDocumentPassingOnAtlasProject: false, + enableGenAIFeaturesAtlasOrg: false, enableMultipleConnectionSystem: true, enablePerformanceAdvisorBanner: true, cloudFeatureRolloutAccess: { @@ -366,9 +370,10 @@ const CompassWeb = ({ onActiveWorkspaceTabChange } onOpenConnectViaModal={onOpenConnectViaModal} - /> + > +