chore(compass-assistant): add confirmation step to explain plan entry point COMPASS-9836

Author: gagik

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,

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<AssistantActionsContextType>({
-    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);
     },
   });

packages/compass-assistant/src/assistant-chat.spec.tsx renamed to 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;
+    });
+  });
 });