Skip to content

Commit 01231f9

Browse files
gagikgagik
authored
chore(compass-assistant): hide the assistant buttons when AI features are disabled COMPASS-9760 (#7262)
Matches NLQ behavior when it comes to button visibility by AI feature permissions.
1 parent b2931e1 commit 01231f9

File tree

11 files changed

+222
-109
lines changed

11 files changed

+222
-109
lines changed

‎packages/compass-assistant/src/compass-assistant-drawer.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ export const CompassAssistantDrawer: React.FunctionComponent<{
4949
'data-testid': 'assistant-confirm-clear-chat-modal',
5050
});
5151
if (confirmed) {
52-
clearChat();
52+
clearChat?.();
5353
}
5454
}, [clearChat]);
5555

‎packages/compass-assistant/src/compass-assistant-provider.spec.tsx

Lines changed: 120 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import React from 'react';
22
import {
33
render,
4+
renderHook,
45
screen,
56
userEvent,
67
waitFor,
@@ -10,6 +11,7 @@ import {
1011
import {
1112
AssistantProvider,
1213
CompassAssistantProvider,
14+
useAssistantActions,
1315
type AssistantMessage,
1416
} from './compass-assistant-provider';
1517
import { expect } from 'chai';
@@ -41,6 +43,94 @@ const TestComponent: React.FunctionComponent<{
4143
);
4244
};
4345

46+
describe('useAssistantActions', function () {
47+
const createWrapper = (chat: Chat<AssistantMessage>) => {
48+
function TestWrapper({ children }: { children: React.ReactNode }) {
49+
return (
50+
<DrawerContentProvider>
51+
<AssistantProvider chat={chat}>{children}</AssistantProvider>
52+
</DrawerContentProvider>
53+
);
54+
}
55+
return TestWrapper;
56+
};
57+
58+
it('returns empty object when AI features are disabled via isAIFeatureEnabled', function () {
59+
const { result } = renderHook(() => useAssistantActions(), {
60+
wrapper: createWrapper(createMockChat({ messages: [] })),
61+
preferences: {
62+
enableAIAssistant: true,
63+
// These control isAIFeatureEnabled
64+
enableGenAIFeatures: false,
65+
enableGenAIFeaturesAtlasOrg: true,
66+
cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
67+
},
68+
});
69+
70+
expect(result.current).to.deep.equal({});
71+
});
72+
73+
it('returns empty object when enableGenAIFeaturesAtlasOrg is disabled', function () {
74+
const { result } = renderHook(() => useAssistantActions(), {
75+
wrapper: createWrapper(createMockChat({ messages: [] })),
76+
preferences: {
77+
enableAIAssistant: true,
78+
enableGenAIFeatures: true,
79+
enableGenAIFeaturesAtlasOrg: false,
80+
cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
81+
},
82+
});
83+
84+
expect(result.current).to.deep.equal({});
85+
});
86+
87+
it('returns empty object when cloudFeatureRolloutAccess is disabled', function () {
88+
const { result } = renderHook(() => useAssistantActions(), {
89+
wrapper: createWrapper(createMockChat({ messages: [] })),
90+
preferences: {
91+
enableAIAssistant: true,
92+
enableGenAIFeatures: true,
93+
enableGenAIFeaturesAtlasOrg: true,
94+
cloudFeatureRolloutAccess: { GEN_AI_COMPASS: false },
95+
},
96+
});
97+
98+
expect(result.current).to.deep.equal({});
99+
});
100+
101+
it('returns empty object when enableAIAssistant preference is disabled', function () {
102+
const { result } = renderHook(() => useAssistantActions(), {
103+
wrapper: createWrapper(createMockChat({ messages: [] })),
104+
preferences: {
105+
enableAIAssistant: false,
106+
enableGenAIFeatures: true,
107+
enableGenAIFeaturesAtlasOrg: true,
108+
cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
109+
},
110+
});
111+
112+
expect(result.current).to.deep.equal({});
113+
});
114+
115+
it('returns actions when both AI features and assistant flag are enabled', function () {
116+
const { result } = renderHook(() => useAssistantActions(), {
117+
wrapper: createWrapper(createMockChat({ messages: [] })),
118+
preferences: {
119+
enableAIAssistant: true,
120+
enableGenAIFeatures: true,
121+
enableGenAIFeaturesAtlasOrg: true,
122+
cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
123+
},
124+
});
125+
126+
expect(Object.keys(result.current)).to.have.length.greaterThan(0);
127+
expect(result.current.interpretExplainPlan).to.be.a('function');
128+
expect(result.current.interpretConnectionError).to.be.a('function');
129+
expect(result.current.tellMoreAboutInsight).to.be.a('function');
130+
expect(result.current.clearChat).to.be.a('function');
131+
});
132+
});
133+
44134
describe('AssistantProvider', function () {
45135
const mockMessages: AssistantMessage[] = [
46136
{
@@ -57,15 +147,25 @@ describe('AssistantProvider', function () {
57147

58148
it('always renders children', function () {
59149
render(<TestComponent chat={createMockChat({ messages: [] })} />, {
60-
preferences: { enableAIAssistant: true },
150+
preferences: {
151+
enableAIAssistant: true,
152+
enableGenAIFeatures: true,
153+
enableGenAIFeaturesAtlasOrg: true,
154+
cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
155+
},
61156
});
62157

63158
expect(screen.getByTestId('provider-children')).to.exist;
64159
});
65160

66161
it('does not render assistant drawer when AI assistant is disabled', function () {
67162
render(<TestComponent chat={createMockChat({ messages: [] })} />, {
68-
preferences: { enableAIAssistant: false },
163+
preferences: {
164+
enableAIAssistant: false,
165+
enableGenAIFeatures: true,
166+
enableGenAIFeaturesAtlasOrg: true,
167+
cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
168+
},
69169
});
70170

71171
expect(screen.getByTestId('provider-children')).to.exist;
@@ -75,7 +175,12 @@ describe('AssistantProvider', function () {
75175

76176
it('renders the assistant drawer as the first drawer item when AI assistant is enabled', function () {
77177
render(<TestComponent chat={createMockChat({ messages: [] })} />, {
78-
preferences: { enableAIAssistant: true },
178+
preferences: {
179+
enableAIAssistant: true,
180+
enableGenAIFeatures: true,
181+
enableGenAIFeaturesAtlasOrg: true,
182+
cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
183+
},
79184
});
80185

81186
expect(screen.getByTestId('lg-drawer-toolbar-icon_button-0')).to.have.attr(
@@ -96,7 +201,12 @@ describe('AssistantProvider', function () {
96201
mockChat: Chat<AssistantMessage>
97202
): Promise<ReturnType<typeof render>> {
98203
const result = render(<TestComponent chat={mockChat} autoOpen={true} />, {
99-
preferences: { enableAIAssistant: true },
204+
preferences: {
205+
enableAIAssistant: true,
206+
enableGenAIFeatures: true,
207+
enableGenAIFeaturesAtlasOrg: true,
208+
cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
209+
},
100210
});
101211

102212
await waitFor(() => {
@@ -281,7 +391,12 @@ describe('AssistantProvider', function () {
281391
<MockedProvider chat={new Chat({})} />
282392
</DrawerContentProvider>,
283393
{
284-
preferences: { enableAIAssistant: true },
394+
preferences: {
395+
enableAIAssistant: true,
396+
enableGenAIFeatures: true,
397+
enableGenAIFeaturesAtlasOrg: true,
398+
cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
399+
},
285400
}
286401
);
287402

‎packages/compass-assistant/src/compass-assistant-provider.tsx

Lines changed: 19 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,10 @@ import {
1616
type EntryPointMessage,
1717
type ProactiveInsightsContext,
1818
} from './prompts';
19-
import { usePreference } from 'compass-preferences-model/provider';
19+
import {
20+
useIsAIFeatureEnabled,
21+
usePreference,
22+
} from 'compass-preferences-model/provider';
2023
import { createLoggerLocator } from '@mongodb-js/compass-logging/provider';
2124
import type { ConnectionInfo } from '@mongodb-js/connection-info';
2225
import { useTelemetry } from '@mongodb-js/compass-telemetry/provider';
@@ -37,22 +40,22 @@ export const AssistantContext = createContext<AssistantContextType | null>(
3740
);
3841

3942
type AssistantActionsContextType = {
40-
interpretExplainPlan: ({
43+
interpretExplainPlan?: ({
4144
namespace,
4245
explainPlan,
4346
}: {
4447
namespace: string;
4548
explainPlan: string;
4649
}) => void;
47-
interpretConnectionError: ({
50+
interpretConnectionError?: ({
4851
connectionInfo,
4952
error,
5053
}: {
5154
connectionInfo: ConnectionInfo;
5255
error: Error;
5356
}) => void;
54-
clearChat: () => void;
55-
tellMoreAboutInsight: (context: ProactiveInsightsContext) => void;
57+
clearChat?: () => void;
58+
tellMoreAboutInsight?: (context: ProactiveInsightsContext) => void;
5659
};
5760
export const AssistantActionsContext =
5861
createContext<AssistantActionsContextType>({
@@ -62,29 +65,21 @@ export const AssistantActionsContext =
6265
clearChat: () => {},
6366
});
6467

65-
export function useAssistantActions(): AssistantActionsContextType & {
66-
isAssistantEnabled: boolean;
67-
} {
68-
const isAssistantEnabled = usePreference('enableAIAssistant');
68+
export function useAssistantActions(): AssistantActionsContextType {
69+
const isAIFeatureEnabled = useIsAIFeatureEnabled();
70+
const isAssistantFlagEnabled = usePreference('enableAIAssistant');
71+
const actions = useContext(AssistantActionsContext);
72+
if (!isAIFeatureEnabled || !isAssistantFlagEnabled) {
73+
return {};
74+
}
6975

70-
return {
71-
...useContext(AssistantActionsContext),
72-
isAssistantEnabled,
73-
};
76+
return actions;
7477
}
7578

7679
export const compassAssistantServiceLocator = createServiceLocator(function () {
77-
const { isAssistantEnabled, ...actions } = useAssistantActions();
80+
const actions = useAssistantActions();
7881

79-
const assistantEnabledRef = useRef(isAssistantEnabled);
80-
assistantEnabledRef.current = isAssistantEnabled;
81-
82-
return {
83-
...actions,
84-
getIsAssistantEnabled() {
85-
return assistantEnabledRef.current;
86-
},
87-
};
82+
return actions;
8883
}, 'compassAssistantLocator');
8984

9085
export type CompassAssistantService = ReturnType<
@@ -160,7 +155,7 @@ export const CompassAssistantProvider = registerCompassPlugin(
160155
transport: new DocsProviderTransport({
161156
baseUrl: atlasService.assistantApiEndpoint(),
162157
}),
163-
onError: (err) => {
158+
onError: (err: Error) => {
164159
logger.log.error(
165160
logger.mongoLogId(1_001_000_370),
166161
'Assistant',

‎packages/compass-connections/src/components/connection-status-notifications.tsx

Lines changed: 19 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,8 @@ export function getConnectingStatusText(connectionInfo: ConnectionInfo) {
4141
type ConnectionErrorToastBodyProps = {
4242
info?: ConnectionInfo | null;
4343
error: Error;
44-
showReviewButton: boolean;
45-
showDebugButton: boolean;
46-
onReview: () => void;
47-
onDebug: () => void;
44+
onReview?: () => void;
45+
onDebug?: () => void;
4846
};
4947

5048
const connectionErrorToastStyles = css({
@@ -94,8 +92,6 @@ const debugActionStyles = css({
9492
function ConnectionErrorToastBody({
9593
info,
9694
error,
97-
showReviewButton,
98-
showDebugButton,
9995
onReview,
10096
onDebug,
10197
}: ConnectionErrorToastBodyProps): React.ReactElement {
@@ -111,7 +107,7 @@ function ConnectionErrorToastBody({
111107
<span data-testid="connection-error-text">{error.message}</span>
112108
</span>
113109
<span className={connectionErrorActionsStyles}>
114-
{info && showReviewButton && (
110+
{info && onReview && (
115111
<span>
116112
<Button
117113
onClick={onReview}
@@ -122,7 +118,7 @@ function ConnectionErrorToastBody({
122118
</Button>
123119
</span>
124120
)}
125-
{info && showDebugButton && (
121+
{info && onDebug && (
126122
<span className={debugActionStyles}>
127123
<Icon glyph="Sparkle" size="small"></Icon>
128124
<Link
@@ -182,8 +178,6 @@ const openConnectionSucceededToast = (connectionInfo: ConnectionInfo) => {
182178
const openConnectionFailedToast = ({
183179
connectionInfo,
184180
error,
185-
showReviewButton,
186-
showDebugButton,
187181
onReviewClick,
188182
onDebugClick,
189183
}: {
@@ -192,10 +186,8 @@ const openConnectionFailedToast = ({
192186
// can happen is autoconnect flow
193187
connectionInfo: ConnectionInfo | null | undefined;
194188
error: Error;
195-
showReviewButton: boolean;
196-
showDebugButton: boolean;
197-
onReviewClick: () => void;
198-
onDebugClick: () => void;
189+
onReviewClick?: () => void;
190+
onDebugClick?: () => void;
199191
}) => {
200192
const failedToastId = connectionInfo?.id ?? 'failed';
201193

@@ -206,19 +198,19 @@ const openConnectionFailedToast = ({
206198
<ConnectionErrorToastBody
207199
info={connectionInfo}
208200
error={error}
209-
showReviewButton={showReviewButton}
210-
showDebugButton={showDebugButton}
211-
onReview={() => {
212-
if (!showDebugButton) {
213-
// don't close the toast if there are two actions so that the user
214-
// can still use the other one
215-
closeToast(`connection-status--${failedToastId}`);
216-
}
217-
onReviewClick();
218-
}}
219-
onDebug={() => {
220-
onDebugClick();
221-
}}
201+
onReview={
202+
onReviewClick
203+
? () => {
204+
if (!onDebugClick) {
205+
// don't close the toast if there are two actions so that the user
206+
// can still use the other one
207+
closeToast(`connection-status--${failedToastId}`);
208+
}
209+
onReviewClick();
210+
}
211+
: undefined
212+
}
213+
onDebug={onDebugClick}
222214
/>
223215
),
224216
variant: 'warning',

‎packages/compass-connections/src/stores/connections-store-redux.spec.tsx

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,9 @@ describe('CompassConnections store', function () {
154154
const { connectionsStore } = renderCompassConnections({
155155
preferences: {
156156
enableAIAssistant: true,
157+
enableGenAIFeatures: true,
158+
enableGenAIFeaturesAtlasOrg: true,
159+
cloudFeatureRolloutAccess: { GEN_AI_COMPASS: true },
157160
},
158161
connectFn: sinon
159162
.stub()

0 commit comments

Comments
 (0)