diff --git a/packages/compass-e2e-tests/helpers/assistant-service.ts b/packages/compass-e2e-tests/helpers/assistant-service.ts new file mode 100644 index 00000000000..4504b78f1c1 --- /dev/null +++ b/packages/compass-e2e-tests/helpers/assistant-service.ts @@ -0,0 +1,263 @@ +import http from 'http'; +import { once } from 'events'; +import type { AddressInfo } from 'net'; + +export type MockAssistantResponse = { + status: number; + body: string; +}; + +function sendStreamingResponse(res: http.ServerResponse, content: string) { + // OpenAI Responses API streaming response format using Server-Sent Events + res.writeHead(200, { + 'Content-Type': 'text/event-stream; charset=utf-8', + 'Cache-Control': 'no-cache', + Connection: 'keep-alive', + 'Transfer-Encoding': 'chunked', + }); + + const responseId = `resp_${Date.now()}`; + const itemId = `item_${Date.now()}`; + let sequenceNumber = 0; + + // Send response.created event + res.write( + `data: ${JSON.stringify({ + type: 'response.created', + response: { + id: responseId, + object: 'realtime.response', + status: 'in_progress', + output: [], + usage: { + input_tokens: 0, + output_tokens: 0, + total_tokens: 0, + }, + }, + sequence_number: sequenceNumber++, + })}\n\n` + ); + + // Send output_item.added event + res.write( + `data: ${JSON.stringify({ + type: 'response.output_item.added', + response_id: responseId, + output_index: 0, + item: { + id: itemId, + object: 'realtime.item', + type: 'message', + role: 'assistant', + content: [], + }, + sequence_number: sequenceNumber++, + })}\n\n` + ); + + // Send the content in chunks + const words = content.split(' '); + let index = 0; + + const sendChunk = () => { + if (index < words.length) { + const word = words[index] + (index < words.length - 1 ? ' ' : ''); + // Send output_text.delta event + res.write( + `data: ${JSON.stringify({ + type: 'response.output_text.delta', + response_id: responseId, + item_id: itemId, + output_index: 0, + delta: word, + sequence_number: sequenceNumber++, + })}\n\n` + ); + index++; + setTimeout(sendChunk, 10); + } else { + // Send output_item.done event + res.write( + `data: ${JSON.stringify({ + type: 'response.output_item.done', + response_id: responseId, + output_index: 0, + item: { + id: itemId, + object: 'realtime.item', + type: 'message', + role: 'assistant', + content: [ + { + type: 'text', + text: content, + }, + ], + }, + sequence_number: sequenceNumber++, + })}\n\n` + ); + + // Send response.completed event + const tokenCount = Math.ceil(content.split(' ').length * 1.3); + res.write( + `data: ${JSON.stringify({ + type: 'response.completed', + response: { + id: responseId, + object: 'realtime.response', + status: 'completed', + output: [ + { + id: itemId, + object: 'realtime.item', + type: 'message', + role: 'assistant', + content: [ + { + type: 'text', + text: content, + }, + ], + }, + ], + usage: { + input_tokens: 10, + output_tokens: tokenCount, + total_tokens: 10 + tokenCount, + }, + }, + sequence_number: sequenceNumber++, + })}\n\n` + ); + + res.write('data: [DONE]\n\n'); + res.end(); + } + }; + + sendChunk(); +} + +export const MOCK_ASSISTANT_SERVER_PORT = 27097; +export async function startMockAssistantServer( + { + response: _response, + }: { + response: MockAssistantResponse; + } = { + response: { + status: 200, + body: 'This is a test response from the AI assistant.', + }, + } +): Promise<{ + clearRequests: () => void; + getResponse: () => MockAssistantResponse; + setResponse: (response: MockAssistantResponse) => void; + getRequests: () => { + content: any; + req: any; + }[]; + endpoint: string; + server: http.Server; + stop: () => Promise; +}> { + let requests: { + content: any; + req: any; + }[] = []; + let response = _response; + const server = http + .createServer((req, res) => { + res.setHeader('Access-Control-Allow-Origin', '*'); + res.setHeader('Access-Control-Allow-Methods', 'POST, OPTIONS'); + res.setHeader( + 'Access-Control-Allow-Headers', + 'Content-Type, Authorization, X-Request-Origin, User-Agent' + ); + + // Handle preflight requests + if (req.method === 'OPTIONS') { + res.writeHead(200); + res.end(); + return; + } + + // Only handle POST requests for chat completions + if (req.method !== 'POST') { + res.writeHead(404); + return res.end('Not Found'); + } + + let body = ''; + req + .setEncoding('utf8') + .on('data', (chunk) => { + body += chunk; + }) + .on('end', () => { + let jsonObject; + try { + jsonObject = JSON.parse(body); + } catch { + res.writeHead(400); + res.setHeader('Content-Type', 'application/json'); + return res.end(JSON.stringify({ error: 'Invalid JSON' })); + } + + requests.push({ + req, + content: jsonObject, + }); + + if (response.status !== 200) { + res.writeHead(response.status); + res.setHeader('Content-Type', 'application/json'); + return res.end(JSON.stringify({ error: response.body })); + } + + // Send streaming response + return sendStreamingResponse(res, response.body); + }); + }) + .listen(MOCK_ASSISTANT_SERVER_PORT); + await once(server, 'listening'); + + // address() returns either a string or AddressInfo. + const address = server.address() as AddressInfo; + + const endpoint = `http://localhost:${address.port}`; + + async function stop() { + server.close(); + await once(server, 'close'); + } + + function clearRequests() { + requests = []; + } + + function getRequests() { + return requests; + } + + function getResponse() { + return response; + } + + function setResponse(newResponse: MockAssistantResponse) { + response = newResponse; + } + + return { + clearRequests, + getRequests, + endpoint, + server, + getResponse, + setResponse, + stop, + }; +} diff --git a/packages/compass-e2e-tests/helpers/selectors.ts b/packages/compass-e2e-tests/helpers/selectors.ts index 5da7673a9f1..54ec3a6db90 100644 --- a/packages/compass-e2e-tests/helpers/selectors.ts +++ b/packages/compass-e2e-tests/helpers/selectors.ts @@ -16,6 +16,9 @@ export const SettingsModalTabSelector = (name: string) => `${SettingsModal} [data-testid="sidebar-${name}-item"]`; export const GeneralSettingsButton = SettingsModalTabSelector('general'); export const GeneralSettingsContent = `${SettingsModal} [data-testid="general-settings"]`; +export const ArtificialIntelligenceSettingsButton = + SettingsModalTabSelector('ai'); +export const ArtificialIntelligenceSettingsContent = `${SettingsModal} [data-testid="gen-ai-settings"]`; export const SettingsInputElement = (settingName: string): string => { return `${SettingsModal} [data-testid="${settingName}"]`; @@ -894,6 +897,9 @@ export const AggregationSavedPipelineCardDeleteButton = ( export const AggregationExplainButton = '[data-testid="pipeline-toolbar-explain-aggregation-button"]'; export const AggregationExplainModal = '[data-testid="explain-plan-modal"]'; +export const ExplainPlanInterpretButton = + '[data-testid="interpret-for-me-button"]'; +export const ExplainPlanCloseButton = '[data-testid="explain-close-button"]'; export const AggregationExplainModalCloseButton = `${AggregationExplainModal} [aria-label*="Close"]`; // Create view from pipeline modal @@ -1510,8 +1516,19 @@ export const SideDrawerCloseButton = `[data-testid="${ }"]`; // Assistant +export const AssistantDrawerButton = 'button[aria-label="MongoDB Assistant"]'; +export const AssistantDrawerCloseButton = `[data-testid="lg-drawer-close_button"]`; export const AssistantChatMessages = '[data-testid="assistant-chat-messages"]'; +export const AssistantChatMessage = '[data-testid^="assistant-message-"]'; +export const AssistantChatInput = '[data-testid="assistant-chat-input"]'; +export const AssistantChatInputTextArea = `${AssistantChatInput} textarea`; +export const AssistantChatSubmitButton = `${AssistantChatInput} button[aria-label="Send message"]`; export const AssistantClearChatButton = '[data-testid="assistant-clear-chat"]'; -export const ConfirmClearChatModal = +export const AssistantConfirmClearChatModal = '[data-testid="assistant-confirm-clear-chat-modal"]'; -export const ConfirmClearChatModalConfirmButton = `${ConfirmClearChatModal} [data-testid="lg-confirmation_modal-footer-confirm_button"]`; +export const AssistantConfirmClearChatModalConfirmButton = `${AssistantConfirmClearChatModal} [data-testid="lg-confirmation_modal-footer-confirm_button"]`; + +// AI Opt-in Modal +export const AIOptInModal = '[data-testid="ai-optin-modal"]'; +export const AIOptInModalAcceptButton = 'button=Use AI Features'; +export const AIOptInModalDeclineLink = 'span=Not now'; diff --git a/packages/compass-e2e-tests/tests/assistant.test.ts b/packages/compass-e2e-tests/tests/assistant.test.ts new file mode 100644 index 00000000000..816e13f4f7f --- /dev/null +++ b/packages/compass-e2e-tests/tests/assistant.test.ts @@ -0,0 +1,571 @@ +import { expect } from 'chai'; + +import type { CompassBrowser } from '../helpers/compass-browser'; +import { startTelemetryServer } from '../helpers/telemetry'; +import type { Telemetry } from '../helpers/telemetry'; +import { + init, + cleanup, + screenshotIfFailed, + DEFAULT_CONNECTION_NAME_1, + skipForWeb, +} from '../helpers/compass'; +import type { Compass } from '../helpers/compass'; +import * as Selectors from '../helpers/selectors'; +import { startMockAtlasServiceServer } from '../helpers/atlas-service'; +import { + MOCK_ASSISTANT_SERVER_PORT, + startMockAssistantServer, +} from '../helpers/assistant-service'; +import type { MockAssistantResponse } from '../helpers/assistant-service'; +import { isTestingWeb } from '../helpers/test-runner-context'; + +if (isTestingWeb()) { + process.env.COMPASS_WEB_FORCE_ENABLE_AI = 'true'; + process.env.COMPASS_ASSISTANT_BASE_URL_OVERRIDE = `http://localhost:${MOCK_ASSISTANT_SERVER_PORT}`; +} + +describe('MongoDB Assistant', function () { + let compass: Compass; + let browser: CompassBrowser; + let telemetry: Telemetry; + + let mockAtlasServer: Awaited>; + let mockAssistantServer: Awaited>; + let sendMessage: ( + text: string, + options?: { + response?: MockAssistantResponse; + expectedResult?: 'success' | 'opt-in'; + } + ) => Promise; + let setAIOptIn: (newValue: boolean) => Promise; + let setAIFeatures: (newValue: boolean) => Promise; + + const testMessage = 'What is MongoDB?'; + const testResponse = 'MongoDB is a database.'; + const dbName = 'test'; + const collectionName = 'entryPoints'; + + before(async function () { + mockAtlasServer = await startMockAtlasServiceServer(); + mockAssistantServer = await startMockAssistantServer(); + + process.env.COMPASS_ATLAS_SERVICE_UNAUTH_BASE_URL_OVERRIDE = + mockAtlasServer.endpoint; + + telemetry = await startTelemetryServer(); + compass = await init(this.test?.fullTitle()); + + sendMessage = async ( + text: string, + { + response = { + status: 200, + body: testResponse, + }, + expectedResult = 'success', + }: { + response?: MockAssistantResponse; + expectedResult?: 'success' | 'opt-in'; + } = {} + ) => { + const existingMessages = await getDisplayedMessages(browser); + + mockAssistantServer.setResponse(response); + const chatInput = browser.$(Selectors.AssistantChatInputTextArea); + await chatInput.waitForDisplayed(); + await chatInput.setValue(text); + const submitButton = browser.$(Selectors.AssistantChatSubmitButton); + await submitButton.click(); + + switch (expectedResult) { + case 'success': + await browser.waitUntil(async () => { + const newMessages = await getDisplayedMessages(browser); + return ( + newMessages.length > existingMessages.length && + newMessages.some( + (message) => + message.text === response?.body && + message.role === 'assistant' + ) && + newMessages.some( + (message) => message.text === text && message.role === 'user' + ) + ); + }); + break; + case 'opt-in': + await browser + .$(Selectors.AIOptInModalAcceptButton) + .waitForDisplayed(); + } + }; + + const setup = async () => { + browser = compass.browser; + await browser.setupDefaultConnections(); + await browser.connectToDefaults(); + await browser.selectConnectionMenuItem( + DEFAULT_CONNECTION_NAME_1, + Selectors.CreateDatabaseButton, + false + ); + await browser.addDatabase(dbName, collectionName); + + await browser.navigateToCollectionTab( + DEFAULT_CONNECTION_NAME_1, + dbName, + collectionName, + 'Aggregations' + ); + }; + + setAIFeatures = async (newValue: boolean) => { + if (!isTestingWeb()) { + const currentValue = await browser.getFeature('enableGenAIFeatures'); + if (currentValue === newValue) { + return; + } + } + + await browser.setFeature('enableGenAIFeatures', newValue); + + if (newValue) { + await browser.$(Selectors.AssistantDrawerButton).waitForDisplayed(); + } else { + await browser.$(Selectors.AssistantDrawerButton).waitForDisplayed({ + reverse: true, + }); + } + }; + + setAIOptIn = async (newValue: boolean) => { + if ( + isTestingWeb() || + ((await browser.getFeature('optInGenAIFeatures')) === true && + newValue === false) + ) { + await cleanup(compass); + // Reseting the opt-in to false can be tricky so it's best to start over in this case. + compass = await init(this.test?.fullTitle(), { firstRun: true }); + await setup(); + return; + } + + await browser.setFeature('optInGenAIFeatures', newValue); + }; + + await setup(); + }); + + after(async function () { + await mockAtlasServer.stop(); + await mockAssistantServer.stop(); + + delete process.env.COMPASS_ATLAS_SERVICE_UNAUTH_BASE_URL_OVERRIDE; + + await cleanup(compass); + await telemetry.stop(); + }); + + afterEach(async function () { + mockAssistantServer.clearRequests(); + await clearChat(browser); + + await screenshotIfFailed(compass, this.currentTest); + }); + + describe('drawer visibility', function () { + it('shows the assistant drawer button when AI features are enabled', async function () { + await setAIFeatures(true); + + const drawerButton = browser.$(Selectors.AssistantDrawerButton); + await drawerButton.waitForDisplayed(); + expect(await drawerButton.isDisplayed()).to.be.true; + }); + + it('does not show the assistant drawer button when AI features are disabled', async function () { + await setAIFeatures(false); + + const drawerButton = browser.$(Selectors.AssistantDrawerButton); + await drawerButton.waitForDisplayed({ reverse: true }); + + await setAIFeatures(true); + }); + + it('can close and open the assistant drawer', async function () { + await openAssistantDrawer(browser); + + await browser.clickVisible(Selectors.AssistantDrawerCloseButton); + + await browser.$(Selectors.AssistantDrawerCloseButton).waitForDisplayed({ + reverse: true, + }); + + await browser.clickVisible(Selectors.AssistantDrawerButton); + + await browser.$(Selectors.AssistantDrawerCloseButton).waitForDisplayed(); + }); + }); + + describe('before opt-in', function () { + before(async function () { + skipForWeb( + this, + 'E2E testing for opt-in on compass-web is not yet implemented' + ); + await setAIOptIn(false); + }); + + it('does not send the message if the user declines the opt-in', async function () { + await openAssistantDrawer(browser); + + await sendMessage(testMessage, { expectedResult: 'opt-in' }); + + await browser.clickVisible(Selectors.AIOptInModalDeclineLink); + + const optInModal = browser.$(Selectors.AIOptInModal); + await optInModal.waitForDisplayed({ reverse: true }); + + const chatInput = browser.$(Selectors.AssistantChatInputTextArea); + expect(await chatInput.getValue()).not.to.equal(testMessage); + + expect(await getDisplayedMessages(browser)).to.deep.equal([]); + + expect(mockAssistantServer.getRequests()).to.be.empty; + }); + + describe('entry points', function () { + it('should display opt-in modal for connection error entry point', async function () { + await browser.connectWithConnectionString( + 'mongodb-invalid://localhost:27017', + { connectionStatus: 'failure' } + ); + await browser.clickVisible( + browser.$(Selectors.ConnectionToastErrorDebugButton) + ); + + const optInModal = browser.$(Selectors.AIOptInModal); + await optInModal.waitForDisplayed(); + expect(await optInModal.isDisplayed()).to.be.true; + + const declineLink = browser.$(Selectors.AIOptInModalDeclineLink); + await declineLink.waitForDisplayed(); + await declineLink.click(); + + await optInModal.waitForDisplayed({ reverse: true }); + + expect(await getDisplayedMessages(browser)).to.deep.equal([]); + }); + + it('should display opt-in modal for explain plan entry point', async function () { + await useExplainPlanEntryPoint(browser); + + const optInModal = browser.$(Selectors.AIOptInModal); + await optInModal.waitForDisplayed(); + expect(await optInModal.isDisplayed()).to.be.true; + + const declineLink = browser.$(Selectors.AIOptInModalDeclineLink); + await declineLink.waitForDisplayed(); + await declineLink.click(); + + await optInModal.waitForDisplayed({ reverse: true }); + + expect(await getDisplayedMessages(browser)).to.deep.equal([]); + }); + }); + }); + + describe('opting in', function () { + before(async function () { + skipForWeb( + this, + 'E2E testing for opt-in on compass-web is not yet implemented' + ); + await setAIOptIn(false); + await openAssistantDrawer(browser); + }); + + it('sends the message if the user opts in', async function () { + await sendMessage(testMessage, { expectedResult: 'opt-in' }); + + const optInModal = browser.$(Selectors.AIOptInModal); + await optInModal.waitForDisplayed(); + expect(await optInModal.isDisplayed()).to.be.true; + + const acceptButton = browser.$(Selectors.AIOptInModalAcceptButton); + await acceptButton.waitForClickable(); + await acceptButton.click(); + + await optInModal.waitForDisplayed({ reverse: true }); + + const chatInput = browser.$(Selectors.AssistantChatInputTextArea); + expect(await chatInput.getValue()).to.equal(''); + + expect(await getDisplayedMessages(browser)).to.deep.equal([ + { text: testMessage, role: 'user' }, + { text: testResponse, role: 'assistant' }, + ]); + }); + }); + + describe('after opt-in', function () { + before(async function () { + await setAIOptIn(true); + await openAssistantDrawer(browser); + }); + + describe('clear chat button', function () { + it('appears only after a message is sent', async function () { + const clearChatButton = browser.$(Selectors.AssistantClearChatButton); + expect(await clearChatButton.isDisplayed()).to.be.false; + }); + + it('should clear the chat when the user clicks the clear chat button', async function () { + await sendMessage(testMessage); + await sendMessage(testMessage); + + expect(await getDisplayedMessages(browser)).to.deep.equal([ + { text: testMessage, role: 'user' }, + { text: testResponse, role: 'assistant' }, + { text: testMessage, role: 'user' }, + { text: testResponse, role: 'assistant' }, + ]); + + await clearChat(browser); + + expect(await getDisplayedMessages(browser)).to.deep.equal([]); + }); + }); + + it('displays multiple messages correctly', async function () { + await sendMessage(testMessage, { + response: { + status: 200, + body: testResponse, + }, + }); + + await sendMessage('This is a different message', { + response: { + status: 200, + body: 'This is a different response', + }, + }); + + expect(await getDisplayedMessages(browser)).to.deep.equal([ + { text: testMessage, role: 'user' }, + { text: testResponse, role: 'assistant' }, + { text: 'This is a different message', role: 'user' }, + { text: 'This is a different response', role: 'assistant' }, + ]); + }); + + it('can copy assistant message to clipboard', async function () { + await sendMessage(testMessage); + + const messageElements = await browser + .$$(Selectors.AssistantChatMessage) + .getElements(); + + const assistantMessage = messageElements[1]; + + const copyButton = assistantMessage.$('[aria-label="Copy message"]'); + await copyButton.waitForDisplayed(); + await copyButton.click(); + + const clipboardText = await browser.execute(() => { + return navigator.clipboard.readText(); + }); + + expect(clipboardText).to.equal(testResponse); + }); + + it('can submit feedback with text', async function () { + await sendMessage(testMessage); + + // Get all message elements + const messageElements = await browser + .$$(Selectors.AssistantChatMessage) + .getElements(); + + const assistantMessage = messageElements[1]; + + const thumbsDownButton = assistantMessage.$( + '[aria-label="Dislike this message"]' + ); + await browser.clickVisible(thumbsDownButton); + + const feedbackTextarea = assistantMessage.$('textarea'); + await feedbackTextarea.waitForDisplayed(); + await feedbackTextarea.setValue('This is a test feedback'); + + await browser.clickVisible(browser.$('button*=Submit')); + + await feedbackTextarea.waitForDisplayed({ reverse: true }); + + const thumbsDownButtonAfter = assistantMessage.$( + '[aria-label="Dislike this message"]' + ); + expect(await thumbsDownButtonAfter.getAttribute('aria-checked')).to.equal( + 'true' + ); + }); + + describe('entry points', function () { + describe('explain plan entry point', function () { + before(async function () { + await setAIOptIn(true); + await setAIFeatures(true); + + mockAssistantServer.setResponse({ + status: 200, + body: 'You should create an index.', + }); + }); + + it('opens assistant with explain plan prompt when clicking "Interpret for me"', async function () { + await useExplainPlanEntryPoint(browser); + + const confirmButton = browser.$('button*=Confirm'); + await confirmButton.waitForDisplayed(); + await confirmButton.click(); + + await browser.waitUntil(async () => { + return (await getDisplayedMessages(browser)).length === 2; + }); + + expect(await getDisplayedMessages(browser)).deep.equal([ + { + text: 'Interpret this explain plan output for me.', + role: 'user', + }, + { text: 'You should create an index.', role: 'assistant' }, + ]); + + expect(mockAssistantServer.getRequests()).to.have.lengthOf(1); + }); + + it('does not send request when user cancels confirmation', async function () { + await useExplainPlanEntryPoint(browser); + + const chatMessages = browser.$(Selectors.AssistantChatMessages); + await chatMessages.waitForDisplayed(); + expect(await chatMessages.getText()).to.include( + 'Please confirm your request' + ); + + // Click Cancel button + const cancelButton = browser.$('button*=Cancel'); + await cancelButton.waitForDisplayed(); + await cancelButton.click(); + + // Wait a bit to ensure no request is sent + await browser.pause(300); + + const finalMessages = await getDisplayedMessages(browser); + expect(finalMessages.length).to.equal(0); + + expect(await chatMessages.getText()).to.include( + 'Please confirm your request' + ); + expect(await chatMessages.getText()).to.include('Request cancelled'); + + // Verify no assistant request was made + expect(mockAssistantServer.getRequests()).to.be.empty; + }); + }); + + describe('error message entry point', function () { + before(function () { + mockAssistantServer.setResponse({ + status: 200, + body: 'You should review the connection string.', + }); + }); + + it('opens assistant with error message view prompt when clicking "Debug for me"', async function () { + await browser.connectWithConnectionString( + 'mongodb-invalid://localhost:27017', + { connectionStatus: 'failure' } + ); + await browser.clickVisible( + browser.$(Selectors.ConnectionToastErrorDebugButton) + ); + + const messages = await getDisplayedMessages(browser); + expect(messages).deep.equal([ + { + text: 'Diagnose why my Compass connection is failing and help me debug it.', + role: 'user', + }, + { + text: 'You should review the connection string.', + role: 'assistant', + }, + ]); + + expect(mockAssistantServer.getRequests()).to.have.lengthOf(1); + }); + }); + }); + }); +}); + +async function openAssistantDrawer(browser: CompassBrowser) { + await browser.clickVisible(Selectors.AssistantDrawerButton); +} + +async function clearChat(browser: CompassBrowser) { + const clearChatButton = browser.$(Selectors.AssistantClearChatButton); + if (await clearChatButton.isDisplayed()) { + await clearChatButton.click(); + const confirmButton = browser.$( + Selectors.AssistantConfirmClearChatModalConfirmButton + ); + await confirmButton.waitForClickable(); + await confirmButton.click(); + } +} + +async function getDisplayedMessages(browser: CompassBrowser) { + await browser.$(Selectors.AssistantChatMessages).waitForDisplayed(); + + const messageElements = await browser + .$$(Selectors.AssistantChatMessage) + .getElements(); + + const displayedMessages = []; + for (const messageElement of messageElements) { + const textElements = await messageElement.$$('p').getElements(); + const isAssistantMessage = + textElements.length !== 1 && + (await textElements[0].getText()) === 'MongoDB Assistant'; + // Get the message text content. + // In case of Assistant messages, skip the MongoDB Assistant text. + const text = isAssistantMessage + ? await textElements[1].getText() + : await textElements[0].getText(); + + displayedMessages.push({ + text: text, + role: isAssistantMessage ? ('assistant' as const) : ('user' as const), + }); + } + + return displayedMessages; +} + +async function useExplainPlanEntryPoint(browser: CompassBrowser) { + await browser.clickVisible(Selectors.AggregationExplainButton); + + await browser.clickVisible(Selectors.ExplainPlanInterpretButton); + + await browser.$(Selectors.AggregationExplainModal).waitForDisplayed({ + reverse: true, + }); + + await browser.$(Selectors.AssistantChatMessages).waitForDisplayed(); +} diff --git a/packages/compass-generative-ai/src/components/ai-optin-modal.tsx b/packages/compass-generative-ai/src/components/ai-optin-modal.tsx index 175be53c93d..2d2989d4628 100644 --- a/packages/compass-generative-ai/src/components/ai-optin-modal.tsx +++ b/packages/compass-generative-ai/src/components/ai-optin-modal.tsx @@ -206,6 +206,7 @@ export const AIOptInModal: React.FunctionComponent = ({ title={`Use AI Features in ${isCloudOptIn ? 'Data Explorer' : 'Compass'}`} open={isOptInModalVisible} onClose={handleModalClose} + data-testid="ai-optin-modal" // TODO Button Disabling className={!isProjectAIEnabled ? currentDisabledButtonStyles : undefined} buttonText="Use AI Features" diff --git a/packages/compass-web/sandbox/index.tsx b/packages/compass-web/sandbox/index.tsx index 897d284d6c4..a5a6848272f 100644 --- a/packages/compass-web/sandbox/index.tsx +++ b/packages/compass-web/sandbox/index.tsx @@ -37,6 +37,9 @@ function getMetaEl(name: string) { const App = () => { const [currentTab, updateCurrentTab] = useWorkspaceTabRouter(); const { status, projectParams } = useAtlasProxySignIn(); + + const forceEnableAI = process.env.COMPASS_WEB_FORCE_ENABLE_AI === 'true'; + const { projectId, csrfToken, @@ -111,12 +114,17 @@ const App = () => { enableRollingIndexes: isAtlas, showDisabledConnections: true, enableGenAIFeaturesAtlasProject: - isAtlas && !!enableGenAIFeaturesAtlasProject, + forceEnableAI || (isAtlas && !!enableGenAIFeaturesAtlasProject), enableGenAISampleDocumentPassing: - isAtlas && !!enableGenAISampleDocumentPassing, + forceEnableAI || + (isAtlas && !!enableGenAISampleDocumentPassing), enableGenAIFeaturesAtlasOrg: - isAtlas && !!enableGenAIFeaturesAtlasOrg, - optInGenAIFeatures: isAtlas && !!optInGenAIFeatures, + forceEnableAI || (isAtlas && !!enableGenAIFeaturesAtlasOrg), + optInGenAIFeatures: + forceEnableAI || (isAtlas && !!optInGenAIFeatures), + cloudFeatureRolloutAccess: forceEnableAI + ? { GEN_AI_COMPASS: true } + : undefined, enableDataModeling: true, enableMyQueries: false, ...groupRolePreferences, diff --git a/packages/compass-web/sandbox/sandbox-atlas-sign-in.tsx b/packages/compass-web/sandbox/sandbox-atlas-sign-in.tsx index 54547da6688..1b04e6e1bc4 100644 --- a/packages/compass-web/sandbox/sandbox-atlas-sign-in.tsx +++ b/packages/compass-web/sandbox/sandbox-atlas-sign-in.tsx @@ -128,19 +128,25 @@ export function useAtlasProxySignIn(): AtlasLoginReturnValue { featureFlags: { groupEnabledFeatureFlags }, userRoles, } = params; + const forceEnableAI = + process.env.COMPASS_WEB_FORCE_ENABLE_AI === 'true'; setProjectParams({ projectId, csrfToken, csrfTime, - optInGenAIFeatures: isOptedIntoDataExplorerGenAIFeatures, - enableGenAIFeaturesAtlasOrg: genAIFeaturesEnabled, + optInGenAIFeatures: + forceEnableAI || isOptedIntoDataExplorerGenAIFeatures, + enableGenAIFeaturesAtlasOrg: forceEnableAI || genAIFeaturesEnabled, enableGenAISampleDocumentPassing: + forceEnableAI || !groupEnabledFeatureFlags.includes( 'DISABLE_DATA_EXPLORER_GEN_AI_SAMPLE_DOCUMENT_PASSING' ), - enableGenAIFeaturesAtlasProject: groupEnabledFeatureFlags.includes( - 'ENABLE_DATA_EXPLORER_GEN_AI_FEATURES' - ), + enableGenAIFeaturesAtlasProject: + forceEnableAI || + groupEnabledFeatureFlags.includes( + 'ENABLE_DATA_EXPLORER_GEN_AI_FEATURES' + ), userRoles, }); setStatus('signed-in'); diff --git a/packages/compass-web/webpack.config.js b/packages/compass-web/webpack.config.js index d4b43ca855d..2ae15bdc4b3 100644 --- a/packages/compass-web/webpack.config.js +++ b/packages/compass-web/webpack.config.js @@ -212,6 +212,16 @@ module.exports = (env, args) => { // Can be either `web` or `webdriverio`, helpful if we need special // behavior for tests in sandbox 'process.env.APP_ENV': JSON.stringify(process.env.APP_ENV ?? 'web'), + 'process.env.COMPASS_WEB_FORCE_ENABLE_AI': JSON.stringify( + process.env.COMPASS_WEB_FORCE_ENABLE_AI ?? 'false' + ), + ...(process.env.COMPASS_ASSISTANT_BASE_URL_OVERRIDE + ? { + 'process.env.COMPASS_ASSISTANT_BASE_URL_OVERRIDE': JSON.stringify( + process.env.COMPASS_ASSISTANT_BASE_URL_OVERRIDE + ), + } + : {}), }), new webpack.ProvidePlugin({