diff --git a/packages/compass-assistant/src/compass-assistant-provider.spec.tsx b/packages/compass-assistant/src/compass-assistant-provider.spec.tsx index c81f43ecf67..58351688d07 100644 --- a/packages/compass-assistant/src/compass-assistant-provider.spec.tsx +++ b/packages/compass-assistant/src/compass-assistant-provider.spec.tsx @@ -10,6 +10,7 @@ import { } from '@mongodb-js/testing-library-compass'; import { CompassAssistantProvider, + createDefaultChat, useAssistantActions, type AssistantMessage, } from './compass-assistant-provider'; @@ -24,8 +25,10 @@ import { import type { AtlasAuthService } from '@mongodb-js/atlas-service/provider'; import type { AtlasService } from '@mongodb-js/atlas-service/provider'; import { CompassAssistantDrawer } from './compass-assistant-drawer'; -import { createMockChat } from '../test/utils'; +import { createBrokenTransport, createMockChat } from '../test/utils'; import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider'; +import type { TrackFunction } from '@mongodb-js/compass-telemetry'; +import { createLogger } from '@mongodb-js/compass-logging'; function createMockProvider({ mockAtlasService, @@ -463,6 +466,44 @@ describe('CompassAssistantProvider', function () { expect(screen.queryByText('Hello assistant!')).to.not.exist; }); + describe('error handling with default chat', function () { + it('fires a telemetry event and displays error banner when error occurs', async function () { + const track = sinon.stub(); + const chat = createDefaultChat({ + options: { + transport: createBrokenTransport(), + }, + originForPrompt: 'mongodb-compass', + appNameForPrompt: 'MongoDB Compass', + atlasService: { + assistantApiEndpoint: sinon + .stub() + .returns('https://localhost:3000'), + } as unknown as AtlasService, + logger: createLogger('COMPASS-ASSISTANT-TEST'), + track: track as unknown as TrackFunction, + }); + await renderOpenAssistantDrawer({ + chat, + }); + + // Send a message + userEvent.type( + screen.getByPlaceholderText('Ask a question'), + 'Hello assistant!' + ); + userEvent.click(screen.getByLabelText('Send message')); + + await waitFor(() => { + expect(screen.getByText(/Test connection error/)).to.exist; + }); + + expect(track).to.have.been.calledWith('Assistant Response Failed', { + error_name: 'ConnectionError', + }); + }); + }); + describe('clear chat button', function () { it('is hidden when the chat is empty', async function () { const mockChat = createMockChat({ messages: [] }); @@ -613,49 +654,47 @@ describe('CompassAssistantProvider', function () { }); }); - describe('CompassAssistantProvider', function () { - it('uses the Atlas Service assistantApiEndpoint', async function () { - const mockAtlasService = { - assistantApiEndpoint: sinon - .stub() - .returns('https://example.com/assistant/api/v1'), - }; + it('uses the Atlas Service assistantApiEndpoint', async function () { + const mockAtlasService = { + assistantApiEndpoint: sinon + .stub() + .returns('https://example.com/assistant/api/v1'), + }; - const mockAtlasAiService = { - ensureAiFeatureAccess: sinon.stub().callsFake(() => { - return Promise.resolve(); - }), - }; + const mockAtlasAiService = { + ensureAiFeatureAccess: sinon.stub().callsFake(() => { + return Promise.resolve(); + }), + }; - const mockAtlasAuthService = {}; + const mockAtlasAuthService = {}; - const MockedProvider = CompassAssistantProvider.withMockServices({ - atlasService: mockAtlasService as unknown as AtlasService, - atlasAiService: mockAtlasAiService as unknown as AtlasAiService, - atlasAuthService: mockAtlasAuthService as unknown as AtlasAuthService, - }); + const MockedProvider = CompassAssistantProvider.withMockServices({ + atlasService: mockAtlasService as unknown as AtlasService, + atlasAiService: mockAtlasAiService as unknown as AtlasAiService, + atlasAuthService: mockAtlasAuthService as unknown as AtlasAuthService, + }); - render( - - - - , - { - preferences: { - enableAIAssistant: true, - enableGenAIFeatures: true, - enableGenAIFeaturesAtlasOrg: true, - cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true }, - }, - } - ); + render( + + + + , + { + preferences: { + enableAIAssistant: true, + enableGenAIFeatures: true, + enableGenAIFeaturesAtlasOrg: true, + cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true }, + }, + } + ); - await waitFor(() => { - expect(mockAtlasService.assistantApiEndpoint.calledOnce).to.be.true; - }); + await waitFor(() => { + expect(mockAtlasService.assistantApiEndpoint.calledOnce).to.be.true; }); }); }); diff --git a/packages/compass-assistant/src/compass-assistant-provider.tsx b/packages/compass-assistant/src/compass-assistant-provider.tsx index 8fd88a8dd7f..e9ced74b509 100644 --- a/packages/compass-assistant/src/compass-assistant-provider.tsx +++ b/packages/compass-assistant/src/compass-assistant-provider.tsx @@ -8,6 +8,7 @@ import { } from '@mongodb-js/compass-app-registry'; import { atlasAuthServiceLocator, + type AtlasService, atlasServiceLocator, } from '@mongodb-js/atlas-service/provider'; import { DocsProviderTransport } from './docs-provider-transport'; @@ -23,9 +24,16 @@ import { useIsAIFeatureEnabled, usePreference, } from 'compass-preferences-model/provider'; -import { createLoggerLocator } from '@mongodb-js/compass-logging/provider'; +import { + createLoggerLocator, + type Logger, +} from '@mongodb-js/compass-logging/provider'; import type { ConnectionInfo } from '@mongodb-js/connection-info'; -import { useTelemetry } from '@mongodb-js/compass-telemetry/provider'; +import { + telemetryLocator, + type TrackFunction, + 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'; @@ -280,29 +288,20 @@ export const CompassAssistantProvider = registerCompassPlugin( ); }, - activate: (initialProps, { atlasService, atlasAiService, logger }) => { + activate: ( + { chat: initialChat, originForPrompt, appNameForPrompt }, + { atlasService, atlasAiService, logger, track } + ) => { const chat = - initialProps.chat ?? - new Chat({ - transport: new DocsProviderTransport({ - origin: initialProps.originForPrompt, - instructions: buildConversationInstructionsPrompt({ - target: initialProps.appNameForPrompt, - }), - model: createOpenAI({ - baseURL: atlasService.assistantApiEndpoint(), - apiKey: '', - }).responses('mongodb-chat-latest'), - }), - onError: (err: Error) => { - logger.log.error( - logger.mongoLogId(1_001_000_370), - 'Assistant', - 'Failed to send a message', - { err } - ); - }, + initialChat ?? + createDefaultChat({ + originForPrompt, + appNameForPrompt, + atlasService, + logger, + track, }); + return { store: { state: { chat, atlasAiService } }, deactivate: () => {}, @@ -313,6 +312,51 @@ export const CompassAssistantProvider = registerCompassPlugin( atlasService: atlasServiceLocator, atlasAiService: atlasAiServiceLocator, atlasAuthService: atlasAuthServiceLocator, + track: telemetryLocator, logger: createLoggerLocator('COMPASS-ASSISTANT'), } ); + +export function createDefaultChat({ + originForPrompt, + appNameForPrompt, + atlasService, + logger, + track, + options, +}: { + originForPrompt: string; + appNameForPrompt: string; + atlasService: AtlasService; + logger: Logger; + track: TrackFunction; + options?: { + transport: Chat['transport']; + }; +}): Chat { + return new Chat({ + transport: + options?.transport ?? + new DocsProviderTransport({ + origin: originForPrompt, + instructions: buildConversationInstructionsPrompt({ + target: appNameForPrompt, + }), + model: createOpenAI({ + baseURL: atlasService.assistantApiEndpoint(), + apiKey: '', + }).responses('mongodb-chat-latest'), + }), + onError: (err: Error) => { + logger.log.error( + logger.mongoLogId(1_001_000_370), + 'Assistant', + 'Failed to send a message', + { err } + ); + track('Assistant Response Failed', { + error_name: err.name, + }); + }, + }); +} diff --git a/packages/compass-assistant/src/components/assistant-chat.spec.tsx b/packages/compass-assistant/src/components/assistant-chat.spec.tsx index b2c2fc54250..98a4d15d6d9 100644 --- a/packages/compass-assistant/src/components/assistant-chat.spec.tsx +++ b/packages/compass-assistant/src/components/assistant-chat.spec.tsx @@ -8,7 +8,7 @@ import { } from '@mongodb-js/testing-library-compass'; import { AssistantChat } from './assistant-chat'; import { expect } from 'chai'; -import { createMockChat } from '../../test/utils'; +import { createBrokenChat, createMockChat } from '../../test/utils'; import type { ConnectionInfo } from '@mongodb-js/connection-info'; import { AssistantActionsContext, @@ -16,6 +16,7 @@ import { } from '../compass-assistant-provider'; import sinon from 'sinon'; import type { TextPart } from 'ai'; +import type { Chat } from '../@ai-sdk/react/chat-react'; describe('AssistantChat', function () { const mockMessages: AssistantMessage[] = [ @@ -46,16 +47,13 @@ describe('AssistantChat', function () { ]; function renderWithChat( - messages: AssistantMessage[], + chat: Chat, { connections, - status, }: { connections?: ConnectionInfo[]; - status?: 'submitted' | 'streaming'; } = {} ) { - const chat = createMockChat({ messages, status }); // The chat component does not use chat.sendMessage() directly, it uses // ensureOptInAndSend() via the AssistantActionsContext. const ensureOptInAndSendStub = sinon @@ -86,7 +84,7 @@ describe('AssistantChat', function () { } it('renders input field and send button', function () { - renderWithChat([]); + renderWithChat(createMockChat({ messages: [] })); const inputField = screen.getByPlaceholderText('Ask a question'); const sendButton = screen.getByLabelText('Send message'); @@ -96,7 +94,7 @@ describe('AssistantChat', function () { }); it('input field accepts text input', function () { - renderWithChat([]); + renderWithChat(createMockChat({ messages: [] })); // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion const inputField = screen.getByPlaceholderText( @@ -109,34 +107,36 @@ describe('AssistantChat', function () { }); it('displays the disclaimer and welcome text', function () { - renderWithChat([]); + renderWithChat(createMockChat({ messages: [] })); expect(screen.getByText(/AI can make mistakes. Review for accuracy./)).to .exist; }); it('displays the welcome text when there are no messages', function () { - renderWithChat([]); + renderWithChat(createMockChat({ messages: [] })); expect(screen.getByText(/Welcome to the MongoDB Assistant!/)).to.exist; }); it('does not display the welcome text when there are messages', function () { - renderWithChat(mockMessages); + renderWithChat(createMockChat({ messages: mockMessages })); expect(screen.queryByText(/Welcome to the MongoDB Assistant!/)).to.not .exist; }); it('displays loading state when chat status is submitted', function () { - renderWithChat([], { status: 'submitted' }); + renderWithChat(createMockChat({ messages: [], status: 'submitted' })); expect(screen.getByText(/MongoDB Assistant is thinking/)).to.exist; }); it('does not display loading in all other cases', function () { - renderWithChat(mockMessages, { status: 'streaming' }); + renderWithChat( + createMockChat({ messages: mockMessages, status: 'streaming' }) + ); expect(screen.queryByText(/MongoDB Assistant is thinking/)).to.not.exist; }); it('send button is disabled when input is empty', function () { - renderWithChat([]); + renderWithChat(createMockChat({ messages: [] })); // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion const sendButton = screen.getByLabelText( @@ -147,7 +147,7 @@ describe('AssistantChat', function () { }); it('send button is enabled when input has text', function () { - renderWithChat([]); + renderWithChat(createMockChat({ messages: [] })); const inputField = screen.getByPlaceholderText('Ask a question'); // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion @@ -161,7 +161,7 @@ describe('AssistantChat', function () { }); it('send button is disabled for whitespace-only input', async function () { - renderWithChat([]); + renderWithChat(createMockChat({ messages: [] })); const inputField = screen.getByPlaceholderText('Ask a question'); // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion @@ -177,7 +177,7 @@ describe('AssistantChat', function () { }); it('displays messages in the chat feed', function () { - renderWithChat(mockMessages); + renderWithChat(createMockChat({ messages: mockMessages })); expect(screen.getByTestId('assistant-message-user')).to.exist; expect(screen.getByTestId('assistant-message-assistant')).to.exist; @@ -240,7 +240,9 @@ describe('AssistantChat', function () { }); it('calls sendMessage when form is submitted', async function () { - const { result, ensureOptInAndSendStub } = renderWithChat([]); + const { result, ensureOptInAndSendStub } = renderWithChat( + createMockChat({ messages: [] }) + ); const { track } = result; const inputField = screen.getByPlaceholderText('Ask a question'); const sendButton = screen.getByLabelText('Send message'); @@ -257,7 +259,7 @@ describe('AssistantChat', function () { }); it('clears input field after successful submission', function () { - renderWithChat([]); + renderWithChat(createMockChat({ messages: [] })); // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion const inputField = screen.getByPlaceholderText( @@ -272,7 +274,9 @@ describe('AssistantChat', function () { }); it('trims whitespace from input before sending', async function () { - const { ensureOptInAndSendStub, result } = renderWithChat([]); + const { ensureOptInAndSendStub, result } = renderWithChat( + createMockChat({ messages: [] }) + ); const { track } = result; const inputField = screen.getByPlaceholderText('Ask a question'); @@ -289,7 +293,9 @@ describe('AssistantChat', function () { }); it('does not call ensureOptInAndSend when input is empty or whitespace-only', function () { - const { ensureOptInAndSendStub } = renderWithChat([]); + const { ensureOptInAndSendStub } = renderWithChat( + createMockChat({ messages: [] }) + ); const inputField = screen.getByPlaceholderText('Ask a question'); const chatForm = screen.getByTestId('assistant-chat-input'); @@ -305,7 +311,7 @@ describe('AssistantChat', function () { }); it('displays user and assistant messages with different styling', function () { - renderWithChat(mockMessages); + renderWithChat(createMockChat({ messages: mockMessages })); const userMessage = screen.getByTestId('assistant-message-user'); const assistantMessage = screen.getByTestId('assistant-message-assistant'); @@ -330,7 +336,7 @@ describe('AssistantChat', function () { }, ]; - renderWithChat(messagesWithMultipleParts); + renderWithChat(createMockChat({ messages: messagesWithMultipleParts })); expect(screen.getByText('Here is part 1. And here is part 2.')).to.exist; }); @@ -349,7 +355,7 @@ describe('AssistantChat', function () { }, ]; - renderWithChat(messagesWithMixedParts); + renderWithChat(createMockChat({ messages: messagesWithMixedParts })); expect(screen.getByText('This is text content. More text content.')).to .exist; @@ -371,7 +377,7 @@ describe('AssistantChat', function () { }, ]; - renderWithChat(messagesWithDisplayText); + renderWithChat(createMockChat({ messages: messagesWithDisplayText })); // Should display the displayText expect( @@ -387,7 +393,7 @@ describe('AssistantChat', function () { describe('feedback buttons', function () { it('shows feedback buttons only for assistant messages', function () { - renderWithChat(mockMessages); + renderWithChat(createMockChat({ messages: mockMessages })); const userMessage = screen.getByTestId('assistant-message-user'); const assistantMessage = screen.getByTestId( @@ -408,7 +414,9 @@ describe('AssistantChat', function () { }); it('tracks positive feedback when thumbs up is clicked', async function () { - const { result } = renderWithChat(mockMessages); + const { result } = renderWithChat( + createMockChat({ messages: mockMessages }) + ); const { track } = result; const assistantMessage = screen.getByTestId( @@ -434,7 +442,9 @@ describe('AssistantChat', function () { }); it('tracks negative feedback when thumbs down is clicked', async function () { - const { result } = renderWithChat(mockMessages); + const { result } = renderWithChat( + createMockChat({ messages: mockMessages }) + ); const { track } = result; const assistantMessage = screen.getByTestId( @@ -461,7 +471,9 @@ describe('AssistantChat', function () { }); it('tracks detailed feedback when feedback text is submitted', async function () { - const { result } = renderWithChat(mockMessages); + const { result } = renderWithChat( + createMockChat({ messages: mockMessages }) + ); const { track } = result; const assistantMessage = screen.getByTestId( @@ -505,15 +517,19 @@ describe('AssistantChat', function () { }); it('tracks it as "chat response" when source is not present', async function () { - const { result } = renderWithChat([ - { - ...mockMessages[1], - metadata: { - ...mockMessages[1].metadata, - source: undefined, - }, - }, - ]); + const { result } = renderWithChat( + createMockChat({ + messages: [ + { + ...mockMessages[1], + metadata: { + ...mockMessages[1].metadata, + source: undefined, + }, + }, + ], + }) + ); const { track } = result; const thumbsDownButton = within( @@ -546,7 +562,7 @@ describe('AssistantChat', function () { }, ]; - renderWithChat(userOnlyMessages); + renderWithChat(createMockChat({ messages: userOnlyMessages })); // Should not find any feedback buttons in the entire component expect(screen.queryByLabelText('Thumbs Up Icon')).to.not.exist; @@ -573,7 +589,7 @@ describe('AssistantChat', function () { }); it('renders confirmation message when message has confirmation metadata', function () { - renderWithChat([mockConfirmationMessage]); + renderWithChat(createMockChat({ messages: [mockConfirmationMessage] })); expect(screen.getByText('Please confirm your request')).to.exist; expect( @@ -584,7 +600,7 @@ describe('AssistantChat', function () { }); it('does not render regular message content when confirmation metadata exists', function () { - renderWithChat([mockConfirmationMessage]); + renderWithChat(createMockChat({ messages: [mockConfirmationMessage] })); // Should not show the message text content when confirmation is present expect(screen.queryByText('This is a confirmation message.')).to.not @@ -592,7 +608,7 @@ describe('AssistantChat', function () { }); it('shows confirmation as pending when it is the last message', function () { - renderWithChat([mockConfirmationMessage]); + renderWithChat(createMockChat({ messages: [mockConfirmationMessage] })); expect(screen.getByText('Confirm')).to.exist; expect(screen.getByText('Cancel')).to.exist; @@ -610,7 +626,7 @@ describe('AssistantChat', function () { }, ]; - renderWithChat(messages); + renderWithChat(createMockChat({ messages: messages })); // The confirmation message (first one) should show as rejected since it's not the last expect(screen.queryByText('Confirm')).to.not.exist; @@ -619,9 +635,9 @@ describe('AssistantChat', function () { }); it('adds new confirmed message when confirmation is confirmed', function () { - const { chat, ensureOptInAndSendStub } = renderWithChat([ - mockConfirmationMessage, - ]); + const { chat, ensureOptInAndSendStub } = renderWithChat( + createMockChat({ messages: [mockConfirmationMessage] }) + ); const confirmButton = screen.getByText('Confirm'); userEvent.click(confirmButton); @@ -638,7 +654,9 @@ describe('AssistantChat', function () { }); it('updates confirmation state to confirmed and adds a new message when confirm button is clicked', function () { - const { chat } = renderWithChat([mockConfirmationMessage]); + const { chat } = renderWithChat( + createMockChat({ messages: [mockConfirmationMessage] }) + ); const confirmButton = screen.getByText('Confirm'); userEvent.click(confirmButton); @@ -657,9 +675,9 @@ describe('AssistantChat', function () { }); 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 { chat, ensureOptInAndSendStub } = renderWithChat( + createMockChat({ messages: [mockConfirmationMessage] }) + ); const cancelButton = screen.getByText('Cancel'); userEvent.click(cancelButton); @@ -678,7 +696,9 @@ describe('AssistantChat', function () { }); it('shows confirmed status after confirmation is confirmed', function () { - const { chat } = renderWithChat([mockConfirmationMessage]); + const { chat } = renderWithChat( + createMockChat({ messages: [mockConfirmationMessage] }) + ); // Verify buttons are initially present expect(screen.getByText('Confirm')).to.exist; @@ -695,7 +715,9 @@ describe('AssistantChat', function () { }); it('shows cancelled status after confirmation is rejected', function () { - const { chat } = renderWithChat([mockConfirmationMessage]); + const { chat } = renderWithChat( + createMockChat({ messages: [mockConfirmationMessage] }) + ); // Verify buttons are initially present expect(screen.getByText('Confirm')).to.exist; @@ -734,7 +756,11 @@ describe('AssistantChat', function () { }, }; - renderWithChat([confirmationMessage1, confirmationMessage2]); + renderWithChat( + createMockChat({ + messages: [confirmationMessage1, confirmationMessage2], + }) + ); expect(screen.getAllByText('Request cancelled')).to.have.length(1); @@ -758,7 +784,9 @@ describe('AssistantChat', function () { }, }; - const { chat } = renderWithChat([messageWithExtraMetadata]); + const { chat } = renderWithChat( + createMockChat({ messages: [messageWithExtraMetadata] }) + ); const confirmButton = screen.getByText('Confirm'); userEvent.click(confirmButton); @@ -777,7 +805,7 @@ describe('AssistantChat', function () { parts: [{ type: 'text', text: 'This is a regular message' }], }; - renderWithChat([regularMessage]); + renderWithChat(createMockChat({ messages: [regularMessage] })); expect(screen.queryByText('Please confirm your request')).to.not.exist; expect(screen.queryByText('Confirm')).to.not.exist; @@ -786,7 +814,9 @@ describe('AssistantChat', function () { }); it('tracks confirmation submitted when confirm button is clicked', async function () { - const { result } = renderWithChat([mockConfirmationMessage]); + const { result } = renderWithChat( + createMockChat({ messages: [mockConfirmationMessage] }) + ); const { track } = result; const confirmButton = screen.getByText('Confirm'); @@ -804,7 +834,9 @@ describe('AssistantChat', function () { }); it('tracks confirmation submitted when cancel button is clicked', async function () { - const { result } = renderWithChat([mockConfirmationMessage]); + const { result } = renderWithChat( + createMockChat({ messages: [mockConfirmationMessage] }) + ); const { track } = result; const cancelButton = screen.getByText('Cancel'); @@ -822,15 +854,19 @@ describe('AssistantChat', function () { }); it('tracks it as "chat response" when source is not present', async function () { - const { result } = renderWithChat([ - { - ...mockConfirmationMessage, - metadata: { - ...mockConfirmationMessage.metadata, - source: undefined, - }, - }, - ]); + const { result } = renderWithChat( + createMockChat({ + messages: [ + { + ...mockConfirmationMessage, + metadata: { + ...mockConfirmationMessage.metadata, + source: undefined, + }, + }, + ], + }) + ); const { track } = result; const confirmButton = screen.getByText('Confirm'); @@ -845,9 +881,55 @@ describe('AssistantChat', function () { }); }); + describe('error handling', function () { + it('displays error banner when error occurs', async function () { + renderWithChat(createBrokenChat()); + + const inputField = screen.getByPlaceholderText('Ask a question'); + const sendButton = screen.getByLabelText('Send message'); + + userEvent.type(inputField, 'What is MongoDB?'); + userEvent.click(sendButton); + + await waitFor(() => { + expect(screen.getByText(/Test connection error. Try clearing the chat/)) + .to.exist; + }); + }); + + it('clears error when close button is clicked', async function () { + const brokenChat = createBrokenChat(); + const clearErrorSpy = sinon.spy(brokenChat, 'clearError'); + + renderWithChat(brokenChat); + + const inputField = screen.getByPlaceholderText('Ask a question'); + const sendButton = screen.getByLabelText('Send message'); + + userEvent.type(inputField, 'What is MongoDB?'); + userEvent.click(sendButton); + + await waitFor(() => { + expect(screen.getByText(/Test connection error. Try clearing the chat/)) + .to.exist; + }); + + const closeButton = screen.getByLabelText('Close Message'); + userEvent.click(closeButton); + + expect(clearErrorSpy).to.have.been.calledOnce; + + await waitFor(() => { + expect( + screen.queryByText(/Test connection error. Try clearing the chat/) + ).to.not.exist; + }); + }); + }); + describe('related sources', function () { it('displays related resources links for assistant messages that include them', async function () { - renderWithChat(mockMessages); + renderWithChat(createMockChat({ messages: mockMessages })); userEvent.click(screen.getByLabelText('Expand Related Resources')); // TODO(COMPASS-9860) can't find the links in test-electron on RHEL and Ubuntu. @@ -868,7 +950,7 @@ describe('AssistantChat', function () { ...message, parts: message.parts.filter((part) => part.type !== 'source-url'), })); - renderWithChat(messages); + renderWithChat(createMockChat({ messages: messages })); expect(screen.queryByLabelText('Expand Related Resources')).to.not.exist; }); @@ -916,7 +998,9 @@ describe('AssistantChat', function () { }, ]; - renderWithChat(messagesWithDuplicateSources); + renderWithChat( + createMockChat({ messages: messagesWithDuplicateSources }) + ); userEvent.click(screen.getByLabelText('Expand Related Resources')); await waitFor(() => { diff --git a/packages/compass-assistant/src/components/assistant-chat.tsx b/packages/compass-assistant/src/components/assistant-chat.tsx index 0f6ee330ce9..5bd5a1ee153 100644 --- a/packages/compass-assistant/src/components/assistant-chat.tsx +++ b/packages/compass-assistant/src/components/assistant-chat.tsx @@ -232,11 +232,6 @@ export const AssistantChat: React.FunctionComponent = ({ const { ensureOptInAndSend } = useContext(AssistantActionsContext); const { messages, status, error, clearError, setMessages } = useChat({ chat, - onError: (error) => { - track('Assistant Response Failed', () => ({ - error_name: error.name, - })); - }, }); const scrollToBottom = useCallback(() => { diff --git a/packages/compass-assistant/test/utils.tsx b/packages/compass-assistant/test/utils.tsx index 211bcc4fd2a..7325ec62726 100644 --- a/packages/compass-assistant/test/utils.tsx +++ b/packages/compass-assistant/test/utils.tsx @@ -5,12 +5,15 @@ import type { AssistantMessage } from '../src/compass-assistant-provider'; export const createMockChat = ({ messages, status, + transport, }: { messages: AssistantMessage[]; status?: 'submitted' | 'streaming'; + transport?: Chat['transport']; }) => { const newChat = new Chat({ messages, + transport, }); sinon.replace(newChat, 'sendMessage', sinon.stub()); if (status) { @@ -20,3 +23,21 @@ export const createMockChat = ({ sendMessage: sinon.SinonStub; }; }; + +export function createBrokenTransport() { + const testError = new Error('Test connection error'); + testError.name = 'ConnectionError'; + const transport = { + sendMessages: sinon.stub().rejects(testError), + reconnectToStream: sinon.stub().resolves(null), + }; + return transport; +} + +export function createBrokenChat() { + const chat = new Chat({ + messages: [], + transport: createBrokenTransport(), + }); + return chat; +}