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;
+}