From 6551ce861dc37ec88ba01b63b29164c4dadc7f8e Mon Sep 17 00:00:00 2001 From: gagik Date: Tue, 16 Sep 2025 13:05:52 +0200 Subject: [PATCH 1/6] chore(compass-assistant): add confirmation step to explain plan entry point COMPASS-9836 --- .../src/compass-assistant-drawer.tsx | 2 +- .../src/compass-assistant-provider.tsx | 29 ++- .../{ => components}/assistant-chat.spec.tsx | 236 +++++++++++++++++- .../src/{ => components}/assistant-chat.tsx | 125 +++++++--- .../components/confirmation-message.spec.tsx | 117 +++++++++ .../src/components/confirmation-message.tsx | 104 ++++++++ .../src/docs-provider-transport.ts | 30 ++- 7 files changed, 599 insertions(+), 44 deletions(-) rename packages/compass-assistant/src/{ => components}/assistant-chat.spec.tsx (66%) rename packages/compass-assistant/src/{ => components}/assistant-chat.tsx (68%) create mode 100644 packages/compass-assistant/src/components/confirmation-message.spec.tsx create mode 100644 packages/compass-assistant/src/components/confirmation-message.tsx diff --git a/packages/compass-assistant/src/compass-assistant-drawer.tsx b/packages/compass-assistant/src/compass-assistant-drawer.tsx index de7a37257d3..c61af06e546 100644 --- a/packages/compass-assistant/src/compass-assistant-drawer.tsx +++ b/packages/compass-assistant/src/compass-assistant-drawer.tsx @@ -8,7 +8,7 @@ import { showConfirmation, spacing, } from '@mongodb-js/compass-components'; -import { AssistantChat } from './assistant-chat'; +import { AssistantChat } from './components/assistant-chat'; import { ASSISTANT_DRAWER_ID, AssistantActionsContext, diff --git a/packages/compass-assistant/src/compass-assistant-provider.tsx b/packages/compass-assistant/src/compass-assistant-provider.tsx index b72d95c2878..b2cbb6c9b70 100644 --- a/packages/compass-assistant/src/compass-assistant-provider.tsx +++ b/packages/compass-assistant/src/compass-assistant-provider.tsx @@ -40,6 +40,11 @@ export type AssistantMessage = UIMessage & { * Used for warning messages in cases like using non-genuine MongoDB. */ isPermanent?: boolean; + /** Information for confirmation messages. */ + confirmation?: { + description: string; + state: 'confirmed' | 'rejected' | 'pending'; + }; }; }; @@ -174,7 +179,17 @@ export const AssistantProvider: React.FunctionComponent< const { prompt, displayText } = builder(props); void assistantActionsContext.current.ensureOptInAndSend( - { text: prompt, metadata: { displayText } }, + { + text: prompt, + metadata: { + displayText, + confirmation: { + description: + 'Explain plan metadata, including the original query, may be used to process your request', + state: 'pending', + }, + }, + }, {}, () => { openDrawer(ASSISTANT_DRAWER_ID); @@ -185,17 +200,17 @@ export const AssistantProvider: React.FunctionComponent< } ); }; - }); + }).current; const assistantActionsContext = useRef({ - interpretExplainPlan: createEntryPointHandler.current( + interpretExplainPlan: createEntryPointHandler( 'explain plan', buildExplainPlanPrompt ), - interpretConnectionError: createEntryPointHandler.current( + interpretConnectionError: createEntryPointHandler( 'connection error', buildConnectionErrorPrompt ), - tellMoreAboutInsight: createEntryPointHandler.current( + tellMoreAboutInsight: createEntryPointHandler( 'performance insights', buildProactiveInsightsPrompt ), @@ -220,6 +235,10 @@ export const AssistantProvider: React.FunctionComponent< // place to do tracking. callback(); + if (chat.status === 'streaming') { + await chat.stop(); + } + await chat.sendMessage(message, options); }, }); diff --git a/packages/compass-assistant/src/assistant-chat.spec.tsx b/packages/compass-assistant/src/components/assistant-chat.spec.tsx similarity index 66% rename from packages/compass-assistant/src/assistant-chat.spec.tsx rename to packages/compass-assistant/src/components/assistant-chat.spec.tsx index c4111c8882a..cffc4195a78 100644 --- a/packages/compass-assistant/src/assistant-chat.spec.tsx +++ b/packages/compass-assistant/src/components/assistant-chat.spec.tsx @@ -7,13 +7,14 @@ import { } from '@mongodb-js/testing-library-compass'; import { AssistantChat } from './assistant-chat'; import { expect } from 'chai'; -import { createMockChat } from '../test/utils'; +import { createMockChat } from '../../test/utils'; import type { ConnectionInfo } from '@mongodb-js/connection-info'; import { AssistantActionsContext, type AssistantMessage, -} from './compass-assistant-provider'; +} from '../compass-assistant-provider'; import sinon from 'sinon'; +import type { TextPart } from 'ai'; describe('AssistantChat', function () { const mockMessages: AssistantMessage[] = [ @@ -527,4 +528,235 @@ describe('AssistantChat', function () { expect(screen.queryByLabelText('Thumbs Down Icon')).to.not.exist; }); }); + + describe('messages with confirmation', function () { + let mockConfirmationMessage: AssistantMessage; + + beforeEach(function () { + mockConfirmationMessage = { + id: 'confirmation-test', + role: 'assistant', + parts: [{ type: 'text', text: 'This is a confirmation message.' }], + metadata: { + confirmation: { + state: 'pending', + description: 'Are you sure you want to proceed with this action?', + }, + }, + }; + }); + + it('renders confirmation message when message has confirmation metadata', function () { + renderWithChat([mockConfirmationMessage]); + + expect(screen.getByText('Please confirm your request')).to.exist; + expect( + screen.getByText('Are you sure you want to proceed with this action?') + ).to.exist; + expect(screen.getByText('Confirm')).to.exist; + expect(screen.getByText('Cancel')).to.exist; + }); + + it('does not render regular message content when confirmation metadata exists', function () { + renderWithChat([mockConfirmationMessage]); + + // Should not show the message text content when confirmation is present + expect(screen.queryByText('This is a confirmation message.')).to.not + .exist; + }); + + it('shows confirmation as pending when it is the last message', function () { + renderWithChat([mockConfirmationMessage]); + + expect(screen.getByText('Confirm')).to.exist; + expect(screen.getByText('Cancel')).to.exist; + expect(screen.queryByText('Request confirmed')).to.not.exist; + expect(screen.queryByText('Request cancelled')).to.not.exist; + }); + + it('shows confirmation as rejected when it is not the last message', function () { + const messages: AssistantMessage[] = [ + mockConfirmationMessage, + { + id: 'newer-message', + role: 'user' as const, + parts: [{ type: 'text', text: 'Another message' }], + }, + ]; + + renderWithChat(messages); + + // The confirmation message (first one) should show as rejected since it's not the last + expect(screen.queryByText('Confirm')).to.not.exist; + expect(screen.queryByText('Cancel')).to.not.exist; + expect(screen.getByText('Request cancelled')).to.exist; + }); + + it('adds new confirmed message when confirmation is confirmed', function () { + const { chat, ensureOptInAndSendStub } = renderWithChat([ + mockConfirmationMessage, + ]); + + const confirmButton = screen.getByText('Confirm'); + userEvent.click(confirmButton); + + // Should add a new message without confirmation metadata + expect(chat.messages).to.have.length(2); + const newMessage = chat.messages[1]; + expect(newMessage.id).to.equal('confirmation-test-confirmed'); + expect(newMessage.metadata?.confirmation).to.be.undefined; + expect(newMessage.parts).to.deep.equal(mockConfirmationMessage.parts); + + // Should call ensureOptInAndSend to send the new message + expect(ensureOptInAndSendStub.calledOnce).to.be.true; + }); + + it('updates confirmation state to confirmed and adds a new message when confirm button is clicked', function () { + const { chat } = renderWithChat([mockConfirmationMessage]); + + const confirmButton = screen.getByText('Confirm'); + userEvent.click(confirmButton); + + // Original message should have updated confirmation state + const originalMessage = chat.messages[0]; + expect(originalMessage.metadata?.confirmation?.state).to.equal( + 'confirmed' + ); + + expect(chat.messages).to.have.length(2); + + expect( + screen.getByText((mockConfirmationMessage.parts[0] as TextPart).text) + ).to.exist; + }); + + it('updates confirmation state to rejected and does not add a new message when cancel button is clicked', function () { + const { chat, ensureOptInAndSendStub } = renderWithChat([ + mockConfirmationMessage, + ]); + + const cancelButton = screen.getByText('Cancel'); + userEvent.click(cancelButton); + + // Original message should have updated confirmation state + const originalMessage = chat.messages[0]; + expect(originalMessage.metadata?.confirmation?.state).to.equal( + 'rejected' + ); + + // Should not add a new message + expect(chat.messages).to.have.length(1); + + // Should not call ensureOptInAndSend + expect(ensureOptInAndSendStub.notCalled).to.be.true; + }); + + it('shows confirmed status after confirmation is confirmed', function () { + const { chat } = renderWithChat([mockConfirmationMessage]); + + // Verify buttons are initially present + expect(screen.getByText('Confirm')).to.exist; + expect(screen.getByText('Cancel')).to.exist; + + const confirmButton = screen.getByText('Confirm'); + userEvent.click(confirmButton); + + // The state update should be immediate - check the chat messages + const updatedMessage = chat.messages[0]; + expect(updatedMessage.metadata?.confirmation?.state).to.equal( + 'confirmed' + ); + }); + + it('shows cancelled status after confirmation is rejected', function () { + const { chat } = renderWithChat([mockConfirmationMessage]); + + // Verify buttons are initially present + expect(screen.getByText('Confirm')).to.exist; + expect(screen.getByText('Cancel')).to.exist; + + const cancelButton = screen.getByText('Cancel'); + userEvent.click(cancelButton); + + // The state update should be immediate - check the chat messages + const updatedMessage = chat.messages[0]; + expect(updatedMessage.metadata?.confirmation?.state).to.equal('rejected'); + }); + + it('handles multiple confirmation messages correctly', function () { + const confirmationMessage1: AssistantMessage = { + id: 'confirmation-1', + role: 'assistant', + parts: [{ type: 'text', text: 'First confirmation' }], + metadata: { + confirmation: { + state: 'pending', + description: 'First confirmation description', + }, + }, + }; + + const confirmationMessage2: AssistantMessage = { + id: 'confirmation-2', + role: 'assistant', + parts: [{ type: 'text', text: 'Second confirmation' }], + metadata: { + confirmation: { + state: 'pending', + description: 'Second confirmation description', + }, + }, + }; + + renderWithChat([confirmationMessage1, confirmationMessage2]); + + expect(screen.getAllByText('Request cancelled')).to.have.length(1); + + expect(screen.getAllByText('Confirm')).to.have.length(1); + expect(screen.getAllByText('Cancel')).to.have.length(1); + expect(screen.getByText('Second confirmation description')).to.exist; + }); + + it('preserves other metadata when creating confirmed message', function () { + const messageWithExtraMetadata: AssistantMessage = { + id: 'confirmation-with-metadata', + role: 'assistant', + parts: [{ type: 'text', text: 'Message with extra metadata' }], + metadata: { + confirmation: { + state: 'pending', + description: 'Confirmation description', + }, + displayText: 'Custom display text', + isPermanent: true, + }, + }; + + const { chat } = renderWithChat([messageWithExtraMetadata]); + + const confirmButton = screen.getByText('Confirm'); + userEvent.click(confirmButton); + + // New confirmed message should preserve other metadata + const newMessage = chat.messages[1]; + expect(newMessage.metadata?.displayText).to.equal('Custom display text'); + expect(newMessage.metadata?.isPermanent).to.equal(true); + expect(newMessage.metadata?.confirmation).to.be.undefined; + }); + + it('does not render confirmation component for regular messages', function () { + const regularMessage: AssistantMessage = { + id: 'regular', + role: 'assistant', + parts: [{ type: 'text', text: 'This is a regular message' }], + }; + + renderWithChat([regularMessage]); + + expect(screen.queryByText('Please confirm your request')).to.not.exist; + expect(screen.queryByText('Confirm')).to.not.exist; + expect(screen.queryByText('Cancel')).to.not.exist; + expect(screen.getByText('This is a regular message')).to.exist; + }); + }); }); diff --git a/packages/compass-assistant/src/assistant-chat.tsx b/packages/compass-assistant/src/components/assistant-chat.tsx similarity index 68% rename from packages/compass-assistant/src/assistant-chat.tsx rename to packages/compass-assistant/src/components/assistant-chat.tsx index b727cbc345c..42006308ca5 100644 --- a/packages/compass-assistant/src/assistant-chat.tsx +++ b/packages/compass-assistant/src/components/assistant-chat.tsx @@ -1,8 +1,8 @@ import React, { useCallback, useEffect, useContext } from 'react'; -import type { AssistantMessage } from './compass-assistant-provider'; -import { AssistantActionsContext } from './compass-assistant-provider'; -import type { Chat } from './@ai-sdk/react/chat-react'; -import { useChat } from './@ai-sdk/react/use-chat'; +import type { AssistantMessage } from '../compass-assistant-provider'; +import { AssistantActionsContext } from '../compass-assistant-provider'; +import type { Chat } from '../@ai-sdk/react/chat-react'; +import { useChat } from '../@ai-sdk/react/use-chat'; import { LgChatChatWindow, LgChatLeafygreenChatProvider, @@ -19,8 +19,9 @@ import { LgChatChatDisclaimer, Link, } from '@mongodb-js/compass-components'; +import { ConfirmationMessage } from './confirmation-message'; import { useTelemetry } from '@mongodb-js/compass-telemetry/provider'; -import { NON_GENUINE_WARNING_MESSAGE } from './preset-messages'; +import { NON_GENUINE_WARNING_MESSAGE } from '../preset-messages'; const { DisclaimerText } = LgChatChatDisclaimer; const { ChatWindow } = LgChatChatWindow; @@ -175,18 +176,6 @@ export const AssistantChat: React.FunctionComponent = ({ } }, [hasNonGenuineConnections, chat, setMessages]); - // Transform AI SDK messages to LeafyGreen chat format - const lgMessages = messages.map((message) => ({ - id: message.id, - messageBody: - message.metadata?.displayText || - message.parts - ?.filter((part) => part.type === 'text') - .map((part) => part.text) - .join(''), - isSender: message.role === 'user', - })); - const handleMessageSend = useCallback( (messageBody: string) => { const trimmedMessageBody = messageBody.trim(); @@ -231,6 +220,43 @@ export const AssistantChat: React.FunctionComponent = ({ [track] ); + const handleConfirmation = useCallback( + ( + confirmedMessage: AssistantMessage, + newState: 'confirmed' | 'rejected' + ) => { + setMessages((messages) => { + const newMessages: AssistantMessage[] = messages.map((message) => { + if ( + message.id === confirmedMessage.id && + message.metadata?.confirmation + ) { + message.metadata.confirmation.state = newState; + } + return message; + }); + + // If confirmed, add a new message with the same content but without confirmation metadata + if (newState === 'confirmed') { + newMessages.push({ + ...confirmedMessage, + id: `${confirmedMessage.id}-confirmed`, + metadata: { + ...confirmedMessage.metadata, + confirmation: undefined, + }, + }); + } + return newMessages; + }); + if (newState === 'confirmed') { + // Force the new message request to be sent + void ensureOptInAndSend?.(undefined, {}, () => {}); + } + }, + [ensureOptInAndSend, setMessages] + ); + return (
= ({ className={messageFeedFixesStyles} >
- {lgMessages.map((messageFields) => ( - - {messageFields.isSender === false && ( - { + const { id, role, metadata } = message; + if (metadata?.confirmation) { + const { description, state } = metadata.confirmation; + const isLastMessage = index === messages.length - 1; + + return ( + handleConfirmation(message, 'confirmed')} + onReject={() => handleConfirmation(message, 'rejected')} /> - )} - - ))} + ); + } + + const displayText = + message.metadata?.displayText || + message.parts + ?.filter((part) => part.type === 'text') + .map((part) => part.text) + .join(''); + + const isSender = role === 'user'; + + // Regular message rendering + return ( + + {!isSender && ( + + )} + + ); + })}
This feature is powered by generative AI. See our{' '} @@ -283,7 +344,7 @@ export const AssistantChat: React.FunctionComponent = ({
)} - {lgMessages.length === 0 && ( + {messages.length === 0 && (

Welcome to your MongoDB Assistant.

Ask any question about MongoDB to receive expert guidance and diff --git a/packages/compass-assistant/src/components/confirmation-message.spec.tsx b/packages/compass-assistant/src/components/confirmation-message.spec.tsx new file mode 100644 index 00000000000..b41ff30ee3a --- /dev/null +++ b/packages/compass-assistant/src/components/confirmation-message.spec.tsx @@ -0,0 +1,117 @@ +import React from 'react'; +import { render, screen, userEvent } from '@mongodb-js/testing-library-compass'; +import { ConfirmationMessage } from './confirmation-message'; +import { expect } from 'chai'; +import sinon from 'sinon'; + +describe('ConfirmationMessage', function () { + const defaultProps = { + state: 'pending' as const, + title: 'Test Confirmation', + description: 'Are you sure you want to proceed with this action?', + onConfirm: () => {}, + onReject: () => {}, + }; + + it('renders title and description', function () { + render(); + + expect(screen.getByText(defaultProps.title)).to.exist; + expect(screen.getByText(defaultProps.description)).to.exist; + }); + + describe('pending state', function () { + it('shows confirm and cancel buttons when both handlers are provided', function () { + const onConfirm = sinon.stub(); + const onReject = sinon.stub(); + + render( + + ); + + expect(screen.getByText('Confirm')).to.exist; + expect(screen.getByText('Cancel')).to.exist; + }); + + it('calls onConfirm when confirm button is clicked', function () { + const onConfirm = sinon.stub(); + const onReject = sinon.stub(); + + render( + + ); + + const confirmButton = screen.getByText('Confirm'); + userEvent.click(confirmButton); + + expect(onConfirm.calledOnce).to.be.true; + expect(onReject.notCalled).to.be.true; + }); + + it('calls onReject when cancel button is clicked', function () { + const onConfirm = sinon.stub(); + const onReject = sinon.stub(); + + render( + + ); + + const cancelButton = screen.getByText('Cancel'); + userEvent.click(cancelButton); + + expect(onReject.calledOnce).to.be.true; + expect(onConfirm.notCalled).to.be.true; + }); + + it('does not show status when in pending state', function () { + render( + + ); + + expect(screen.queryByText('Request confirmed')).to.not.exist; + expect(screen.queryByText('Request cancelled')).to.not.exist; + }); + }); + + describe('confirmed and rejected states', function () { + it('shows confirmed status with checkmark icon', function () { + render(); + + expect(screen.getByText('Request confirmed')).to.exist; + // This is sic from the icon library + expect(screen.getByLabelText('Checkmark With Circle Icon')).to.exist; + + expect(screen.queryByText('Confirm')).to.not.exist; + expect(screen.queryByText('Cancel')).to.not.exist; + }); + + it('shows cancelled status', function () { + render(); + + expect(screen.getByText('Request cancelled')).to.exist; + expect(screen.getByLabelText('X With Circle Icon')).to.exist; + + expect(screen.queryByText('Confirm')).to.not.exist; + expect(screen.queryByText('Cancel')).to.not.exist; + }); + }); +}); diff --git a/packages/compass-assistant/src/components/confirmation-message.tsx b/packages/compass-assistant/src/components/confirmation-message.tsx new file mode 100644 index 00000000000..4ff9c77a821 --- /dev/null +++ b/packages/compass-assistant/src/components/confirmation-message.tsx @@ -0,0 +1,104 @@ +import React from 'react'; +import { + Icon, + Body, + Button, + ButtonVariant, + spacing, + css, + palette, + useDarkMode, +} from '@mongodb-js/compass-components'; + +const confirmationMessageStyles = css({ + padding: spacing[300], + borderRadius: spacing[200], + backgroundColor: palette.gray.light3, + border: `1px solid ${palette.gray.light2}`, +}); + +const confirmationStatusStyles = css({ + display: 'flex', + alignItems: 'center', + gap: spacing[100], + marginTop: spacing[200], +}); + +const confirmationStatusTextStyles = css({ + color: palette.gray.dark1, +}); + +const confirmationTitleStyles = css({ + marginBottom: spacing[150], + fontWeight: 600, +}); + +const buttonGroupStyles = css({ + display: 'flex', + gap: spacing[200], + marginTop: spacing[300], + '> button': { + width: '100%', + }, +}); + +interface ConfirmationMessageProps { + state: 'confirmed' | 'rejected' | 'pending'; + title: string; + description: string; + messageBody?: string; + onConfirm: () => void; + onReject: () => void; +} + +export const ConfirmationMessage: React.FunctionComponent< + ConfirmationMessageProps +> = ({ state, title, description, onConfirm, onReject }) => { + const darkMode = useDarkMode(); + + return ( +
+ {title} + + {description} + + {state === 'pending' && ( +
+ + +
+ )} + {state !== 'pending' && ( +
+ + + Request {state === 'confirmed' ? 'confirmed' : 'cancelled'} + +
+ )} +
+ ); +}; diff --git a/packages/compass-assistant/src/docs-provider-transport.ts b/packages/compass-assistant/src/docs-provider-transport.ts index 6db18fd0190..b8d0b42066d 100644 --- a/packages/compass-assistant/src/docs-provider-transport.ts +++ b/packages/compass-assistant/src/docs-provider-transport.ts @@ -1,13 +1,21 @@ import { type ChatTransport, - type UIMessage, type UIMessageChunk, convertToModelMessages, streamText, } from 'ai'; import { createOpenAI } from '@ai-sdk/openai'; +import type { AssistantMessage } from './compass-assistant-provider'; -export class DocsProviderTransport implements ChatTransport { +/** Returns true if the message should be excluded from being sent to the assistant API. */ +export function shouldExcludeMessage({ metadata }: AssistantMessage) { + if (metadata?.confirmation) { + return metadata.confirmation.state !== 'confirmed'; + } + return false; +} + +export class DocsProviderTransport implements ChatTransport { private openai: ReturnType; private instructions: string; @@ -28,10 +36,24 @@ export class DocsProviderTransport implements ChatTransport { sendMessages({ messages, abortSignal, - }: Parameters['sendMessages']>[0]) { + }: Parameters['sendMessages']>[0]) { + const filteredMessages = messages.filter( + (message) => !shouldExcludeMessage(message) + ); + + // If no messages remain after filtering, return an empty stream + if (filteredMessages.length === 0) { + const emptyStream = new ReadableStream({ + start(controller) { + controller.close(); + }, + }); + return Promise.resolve(emptyStream); + } + const result = streamText({ model: this.openai.responses('mongodb-chat-latest'), - messages: convertToModelMessages(messages), + messages: convertToModelMessages(filteredMessages), abortSignal: abortSignal, providerOptions: { openai: { From c50694b4e1ae7a863c9f080d71d99331b9cae284 Mon Sep 17 00:00:00 2001 From: gagik Date: Tue, 16 Sep 2025 13:08:22 +0200 Subject: [PATCH 2/6] chore: prevent re-sending requests --- .../src/docs-provider-transport.ts | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/packages/compass-assistant/src/docs-provider-transport.ts b/packages/compass-assistant/src/docs-provider-transport.ts index b8d0b42066d..68db7bb45f4 100644 --- a/packages/compass-assistant/src/docs-provider-transport.ts +++ b/packages/compass-assistant/src/docs-provider-transport.ts @@ -33,22 +33,30 @@ export class DocsProviderTransport implements ChatTransport { this.instructions = instructions; } + static emptyStream = new ReadableStream({ + start(controller) { + controller.close(); + }, + }); + sendMessages({ messages, abortSignal, }: Parameters['sendMessages']>[0]) { + // If the most recent message is a message that is meant to be excluded + // then we do not need to send this request to the assistant API as it's likely + // redundant otherwise. + if (shouldExcludeMessage(messages[messages.length - 1])) { + return Promise.resolve(DocsProviderTransport.emptyStream); + } + const filteredMessages = messages.filter( (message) => !shouldExcludeMessage(message) ); // If no messages remain after filtering, return an empty stream if (filteredMessages.length === 0) { - const emptyStream = new ReadableStream({ - start(controller) { - controller.close(); - }, - }); - return Promise.resolve(emptyStream); + return Promise.resolve(DocsProviderTransport.emptyStream); } const result = streamText({ From de9f5cfbc4386ce56de319bf1b396b3bda3ae6ed Mon Sep 17 00:00:00 2001 From: gagik Date: Tue, 16 Sep 2025 15:57:45 +0200 Subject: [PATCH 3/6] chore: add tests, refactor docs provider --- .../src/compass-assistant-provider.tsx | 17 +- .../components/confirmation-message.spec.tsx | 5 +- .../src/docs-provider-transport.spec.ts | 202 ++++++++++++++++++ .../src/docs-provider-transport.ts | 17 +- packages/compass-assistant/src/prompts.ts | 38 ++-- 5 files changed, 245 insertions(+), 34 deletions(-) create mode 100644 packages/compass-assistant/src/docs-provider-transport.spec.ts diff --git a/packages/compass-assistant/src/compass-assistant-provider.tsx b/packages/compass-assistant/src/compass-assistant-provider.tsx index b2cbb6c9b70..8bffc1a99c7 100644 --- a/packages/compass-assistant/src/compass-assistant-provider.tsx +++ b/packages/compass-assistant/src/compass-assistant-provider.tsx @@ -29,6 +29,7 @@ import { useTelemetry } from '@mongodb-js/compass-telemetry/provider'; import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider'; import { atlasAiServiceLocator } from '@mongodb-js/compass-generative-ai/provider'; import { buildConversationInstructionsPrompt } from './prompts'; +import { createOpenAI } from '@ai-sdk/openai'; export const ASSISTANT_DRAWER_ID = 'compass-assistant-drawer'; @@ -177,18 +178,11 @@ export const AssistantProvider: React.FunctionComponent< return; } - const { prompt, displayText } = builder(props); + const { prompt, metadata } = builder(props); void assistantActionsContext.current.ensureOptInAndSend( { text: prompt, - metadata: { - displayText, - confirmation: { - description: - 'Explain plan metadata, including the original query, may be used to process your request', - state: 'pending', - }, - }, + metadata, }, {}, () => { @@ -286,10 +280,13 @@ export const CompassAssistantProvider = registerCompassPlugin( initialProps.chat ?? new Chat({ transport: new DocsProviderTransport({ - baseUrl: atlasService.assistantApiEndpoint(), instructions: buildConversationInstructionsPrompt({ target: initialProps.appNameForPrompt, }), + model: createOpenAI({ + baseURL: atlasService.assistantApiEndpoint(), + apiKey: '', + }).responses('mongodb-chat-latest'), }), onError: (err: Error) => { logger.log.error( diff --git a/packages/compass-assistant/src/components/confirmation-message.spec.tsx b/packages/compass-assistant/src/components/confirmation-message.spec.tsx index b41ff30ee3a..379239cf8b4 100644 --- a/packages/compass-assistant/src/components/confirmation-message.spec.tsx +++ b/packages/compass-assistant/src/components/confirmation-message.spec.tsx @@ -97,7 +97,7 @@ describe('ConfirmationMessage', function () { render(); expect(screen.getByText('Request confirmed')).to.exist; - // This is sic from the icon library + // sic from the icon library expect(screen.getByLabelText('Checkmark With Circle Icon')).to.exist; expect(screen.queryByText('Confirm')).to.not.exist; @@ -108,7 +108,8 @@ describe('ConfirmationMessage', function () { render(); expect(screen.getByText('Request cancelled')).to.exist; - expect(screen.getByLabelText('X With Circle Icon')).to.exist; + // sic from the icon library + expect(screen.getByLabelText('XWith Circle Icon')).to.exist; expect(screen.queryByText('Confirm')).to.not.exist; expect(screen.queryByText('Cancel')).to.not.exist; diff --git a/packages/compass-assistant/src/docs-provider-transport.spec.ts b/packages/compass-assistant/src/docs-provider-transport.spec.ts new file mode 100644 index 00000000000..c860e5e671e --- /dev/null +++ b/packages/compass-assistant/src/docs-provider-transport.spec.ts @@ -0,0 +1,202 @@ +import { expect } from 'chai'; +import sinon from 'sinon'; +import { + DocsProviderTransport, + shouldExcludeMessage, +} from './docs-provider-transport'; +import type { AssistantMessage } from './compass-assistant-provider'; +import { MockLanguageModelV2 } from 'ai/test'; +import type { UIMessageChunk } from 'ai'; +import { waitFor } from '@mongodb-js/testing-library-compass'; + +describe.only('DocsProviderTransport', function () { + describe('shouldExcludeMessage', function () { + it('returns false for messages without confirmation metadata', function () { + const message: AssistantMessage = { + id: 'test-1', + role: 'user', + parts: [{ type: 'text', text: 'Hello' }], + }; + + expect(shouldExcludeMessage(message)).to.be.false; + }); + + it('returns true for confirmation messages', function () { + const message: AssistantMessage = { + id: 'test-5', + role: 'assistant', + parts: [{ type: 'text', text: 'Response' }], + metadata: { + confirmation: { + state: 'pending', + description: 'Confirm this action', + }, + }, + }; + + expect(shouldExcludeMessage(message)).to.be.true; + }); + }); + + describe('sending messages', function () { + let mockModel: MockLanguageModelV2; + let doStream: sinon.SinonStub; + let transport: DocsProviderTransport; + let abortController: AbortController; + let sendMessages: ( + params: Partial[0]> + ) => Promise>; + + beforeEach(function () { + // Mock the OpenAI client + doStream = sinon.stub().returns({ + stream: DocsProviderTransport.emptyStream, + request: { + body: { + messages: [], + }, + }, + }); + mockModel = new MockLanguageModelV2({ + doStream, + }); + abortController = new AbortController(); + transport = new DocsProviderTransport({ + instructions: 'Test instructions for MongoDB assistance', + model: mockModel, + }); + sendMessages = (params) => + transport.sendMessages({ + trigger: 'submit-message', + chatId: 'test-chat', + messageId: undefined, + abortSignal: abortController.signal, + messages: [], + ...params, + }); + }); + + afterEach(function () { + sinon.restore(); + }); + + describe('sendMessages', function () { + const userMessage: AssistantMessage = { + id: 'included1', + role: 'user', + parts: [{ type: 'text', text: 'User message' }], + }; + const confirmationPendingMessage: AssistantMessage = { + id: 'test', + role: 'assistant', + parts: [{ type: 'text', text: 'Response' }], + metadata: { + confirmation: { + state: 'pending', + description: 'Confirm this action', + }, + }, + }; + const confirmationConfirmedMessage: AssistantMessage = { + id: 'test', + role: 'assistant', + parts: [{ type: 'text', text: 'Response' }], + metadata: { + confirmation: { + state: 'confirmed', + description: 'Confirmed action', + }, + }, + }; + it('returns empty stream when last message should be excluded', async function () { + const messages: AssistantMessage[] = [ + userMessage, + confirmationConfirmedMessage, + ]; + + const result = await sendMessages({ + messages, + }); + + expect(result).to.equal(DocsProviderTransport.emptyStream); + expect(mockModel.doStreamCalls).to.be.empty; + }); + + it('returns empty stream when all messages are filtered out', async function () { + const messages: AssistantMessage[] = [ + confirmationPendingMessage, + confirmationConfirmedMessage, + ]; + + const result = await sendMessages({ + messages, + }); + + expect(result).to.equal(DocsProviderTransport.emptyStream); + expect(mockModel.doStreamCalls).to.be.empty; + }); + + it('sends filtered messages to AI when valid messages exist', async function () { + await sendMessages({ + messages: [userMessage], + }); + + await waitFor(() => { + expect(doStream).to.have.been.calledOnce; + expect(doStream.firstCall.args[0]).to.deep.include({ + prompt: [ + { + role: 'user', + providerOptions: undefined, + content: [ + { + type: 'text', + text: 'User message', + providerOptions: undefined, + }, + ], + }, + ], + }); + }); + }); + + it('sends only valid messages when confirmation required messages exist', async function () { + await sendMessages({ + messages: [ + confirmationConfirmedMessage, + confirmationPendingMessage, + userMessage, + ], + }); + + await waitFor(() => { + expect(doStream).to.have.been.calledOnce; + expect(doStream.firstCall.args[0]).to.deep.include({ + prompt: [ + { + role: 'user', + providerOptions: undefined, + content: [ + { + type: 'text', + text: 'User message', + providerOptions: undefined, + }, + ], + }, + ], + }); + }); + }); + }); + + // We currently do not support reconnecting to streams but we may want to in the future + describe('reconnectToStream', function () { + it('always returns null', async function () { + const result = await transport.reconnectToStream(); + expect(result).to.be.null; + }); + }); + }); +}); diff --git a/packages/compass-assistant/src/docs-provider-transport.ts b/packages/compass-assistant/src/docs-provider-transport.ts index 68db7bb45f4..68d5d2fc023 100644 --- a/packages/compass-assistant/src/docs-provider-transport.ts +++ b/packages/compass-assistant/src/docs-provider-transport.ts @@ -1,36 +1,33 @@ import { type ChatTransport, + type LanguageModel, type UIMessageChunk, convertToModelMessages, streamText, } from 'ai'; -import { createOpenAI } from '@ai-sdk/openai'; import type { AssistantMessage } from './compass-assistant-provider'; /** Returns true if the message should be excluded from being sent to the assistant API. */ export function shouldExcludeMessage({ metadata }: AssistantMessage) { if (metadata?.confirmation) { - return metadata.confirmation.state !== 'confirmed'; + return true; } return false; } export class DocsProviderTransport implements ChatTransport { - private openai: ReturnType; + private model: LanguageModel; private instructions: string; constructor({ - baseUrl, instructions, + model, }: { - baseUrl: string; instructions: string; + model: LanguageModel; }) { - this.openai = createOpenAI({ - baseURL: baseUrl, - apiKey: '', - }); this.instructions = instructions; + this.model = model; } static emptyStream = new ReadableStream({ @@ -60,7 +57,7 @@ export class DocsProviderTransport implements ChatTransport { } const result = streamText({ - model: this.openai.responses('mongodb-chat-latest'), + model: this.model, messages: convertToModelMessages(filteredMessages), abortSignal: abortSignal, providerOptions: { diff --git a/packages/compass-assistant/src/prompts.ts b/packages/compass-assistant/src/prompts.ts index 8dfaf6dff67..d866ed60efb 100644 --- a/packages/compass-assistant/src/prompts.ts +++ b/packages/compass-assistant/src/prompts.ts @@ -1,5 +1,11 @@ import type { ConnectionInfo } from '@mongodb-js/connection-info'; import { redactConnectionString } from 'mongodb-connection-string-url'; +import type { AssistantMessage } from './compass-assistant-provider'; + +export type EntryPointMessage = { + prompt: string; + metadata: AssistantMessage['metadata']; +}; export const APP_NAMES_FOR_PROMPT = { Compass: 'MongoDB Compass', @@ -36,11 +42,6 @@ You CANNOT: `; }; -export type EntryPointMessage = { - prompt: string; - displayText?: string; -}; - export type ExplainPlanContext = { explainPlan: string; }; @@ -114,7 +115,14 @@ Tell the user if indexes need to be created or modified to enable any recommenda ${explainPlan} `, - displayText: 'Interpret this explain plan output for me.', + metadata: { + displayText: 'Interpret this explain plan output for me.', + confirmation: { + description: + 'Explain plan metadata, including the original query, may be used to process your request', + state: 'pending', + }, + }, }; }; @@ -134,8 +142,6 @@ export const buildProactiveInsightsPrompt = ( switch (context.id) { case 'aggregation-executed-without-index': { return { - displayText: - 'Help me understand the performance impact of running aggregations without an index.', prompt: `The given MongoDB aggregation was executed without an index. Provide a concise human readable explanation that explains why it might degrade performance to not use an index. Please suggest whether an existing index can be used to improve the performance of this query, or if a new index must be created, and describe how it can be accomplished in MongoDB Compass. Do not advise users to create indexes without weighing the pros and cons. @@ -145,12 +151,14 @@ Respond with as much concision and clarity as possible. ${context.stages.join('\n')} `, + metadata: { + displayText: + 'Help me understand the performance impact of running aggregations without an index.', + }, }; } case 'query-executed-without-index': return { - displayText: - 'Help me understand the performance impact of running queries without an index.', prompt: `The given MongoDB query was executed without an index. Provide a concise human readable explanation that explains why it might degrade performance to not use an index. Please suggest whether an existing index can be used to improve the performance of this query, or if a new index must be created, and describe how it can be accomplished in MongoDB Compass. Do not advise users to create indexes without weighing the pros and cons. @@ -160,6 +168,10 @@ Respond with as much concision and clarity as possible. ${context.query} `, + metadata: { + displayText: + 'Help me understand the performance impact of running queries without an index.', + }, }; } }; @@ -188,7 +200,9 @@ ${connectionString} Error message: ${connectionError}`, - displayText: - 'Diagnose why my Compass connection is failing and help me debug it.', + metadata: { + displayText: + 'Diagnose why my Compass connection is failing and help me debug it.', + }, }; }; From 4afb5fadd9964d42bb436d22fb4d592745a3507d Mon Sep 17 00:00:00 2001 From: gagik Date: Tue, 16 Sep 2025 16:03:25 +0200 Subject: [PATCH 4/6] chore: remove only --- packages/compass-assistant/src/docs-provider-transport.spec.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/compass-assistant/src/docs-provider-transport.spec.ts b/packages/compass-assistant/src/docs-provider-transport.spec.ts index c860e5e671e..2143d429971 100644 --- a/packages/compass-assistant/src/docs-provider-transport.spec.ts +++ b/packages/compass-assistant/src/docs-provider-transport.spec.ts @@ -9,7 +9,7 @@ import { MockLanguageModelV2 } from 'ai/test'; import type { UIMessageChunk } from 'ai'; import { waitFor } from '@mongodb-js/testing-library-compass'; -describe.only('DocsProviderTransport', function () { +describe('DocsProviderTransport', function () { describe('shouldExcludeMessage', function () { it('returns false for messages without confirmation metadata', function () { const message: AssistantMessage = { From d5493be56347b918d759c4ed8e434030ea203325 Mon Sep 17 00:00:00 2001 From: Gagik Amaryan Date: Tue, 16 Sep 2025 16:50:11 +0200 Subject: [PATCH 5/6] Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../src/components/assistant-chat.tsx | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/packages/compass-assistant/src/components/assistant-chat.tsx b/packages/compass-assistant/src/components/assistant-chat.tsx index 42006308ca5..a025929d6ff 100644 --- a/packages/compass-assistant/src/components/assistant-chat.tsx +++ b/packages/compass-assistant/src/components/assistant-chat.tsx @@ -231,7 +231,16 @@ export const AssistantChat: React.FunctionComponent = ({ message.id === confirmedMessage.id && message.metadata?.confirmation ) { - message.metadata.confirmation.state = newState; + return { + ...message, + metadata: { + ...message.metadata, + confirmation: { + ...message.metadata.confirmation, + state: newState, + }, + }, + }; } return message; }); From 693c46fc9422382254702efd293a4673e23d5e59 Mon Sep 17 00:00:00 2001 From: gagik Date: Wed, 17 Sep 2025 13:00:03 +0200 Subject: [PATCH 6/6] chore: changes from feedback --- .../src/components/confirmation-message.spec.tsx | 2 +- .../compass-assistant/src/components/confirmation-message.tsx | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/compass-assistant/src/components/confirmation-message.spec.tsx b/packages/compass-assistant/src/components/confirmation-message.spec.tsx index 379239cf8b4..d9b822afa65 100644 --- a/packages/compass-assistant/src/components/confirmation-message.spec.tsx +++ b/packages/compass-assistant/src/components/confirmation-message.spec.tsx @@ -21,7 +21,7 @@ describe('ConfirmationMessage', function () { }); describe('pending state', function () { - it('shows confirm and cancel buttons when both handlers are provided', function () { + it('shows confirm and cancel buttons', function () { const onConfirm = sinon.stub(); const onReject = sinon.stub(); diff --git a/packages/compass-assistant/src/components/confirmation-message.tsx b/packages/compass-assistant/src/components/confirmation-message.tsx index 4ff9c77a821..d6929beeac9 100644 --- a/packages/compass-assistant/src/components/confirmation-message.tsx +++ b/packages/compass-assistant/src/components/confirmation-message.tsx @@ -46,7 +46,6 @@ interface ConfirmationMessageProps { state: 'confirmed' | 'rejected' | 'pending'; title: string; description: string; - messageBody?: string; onConfirm: () => void; onReject: () => void; }