Skip to content

Commit 6c77ed7

Browse files
NickGaganNickGagan
authored
Move selectedMCPServers and selectedTools to Zustand (opendatahub-io#6107)
* Move selectedMCPServers to Zustand * Move selectedTools to Zustand * Move mcpToolSelections after initialValues
1 parent 301cfa0 commit 6c77ed7

13 files changed

Lines changed: 664 additions & 400 deletions

File tree

‎packages/gen-ai/frontend/src/app/Chatbot/ChatbotPlayground.tsx‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ import { isLlamaModelEnabled } from '~/app/utilities';
1919
import { TokenInfo } from '~/app/types';
2020
import useFetchMCPServers from '~/app/hooks/useFetchMCPServers';
2121
import useMCPServerStatuses from '~/app/hooks/useMCPServerStatuses';
22-
import { useMCPToolSelections } from '~/app/hooks/useMCPToolSelections';
2322
import { FILE_UPLOAD_CONFIG, ERROR_MESSAGES, sampleWelcomePrompts } from './const';
2423
import { ChatbotSourceSettingsModal } from './sourceUpload/ChatbotSourceSettingsModal';
2524
import useSourceManagement from './hooks/useSourceManagement';
@@ -35,6 +34,7 @@ import {
3534
selectStreamingEnabled,
3635
selectSelectedModel,
3736
selectConfigIds,
37+
selectSelectedMcpServerIds,
3838
} from './store';
3939
import SourceUploadErrorAlert from './components/alerts/SourceUploadErrorAlert';
4040
import SourceUploadSuccessAlert from './components/alerts/SourceUploadSuccessAlert';
@@ -58,7 +58,6 @@ const ChatbotPlayground: React.FC<ChatbotPlaygroundProps> = ({
5858
}) => {
5959
const { username } = useUserContext();
6060
const { namespace } = React.useContext(GenAiContext);
61-
const { getToolSelections } = useMCPToolSelections();
6261
const { models, modelsLoaded, aiModels, maasModels, lastInput, setLastInput } =
6362
React.useContext(ChatbotContext);
6463

@@ -73,6 +72,8 @@ const ChatbotPlayground: React.FC<ChatbotPlaygroundProps> = ({
7372
const temperature = useChatbotConfigStore(selectTemperature(configId));
7473
const isStreamingEnabled = useChatbotConfigStore(selectStreamingEnabled(configId));
7574
const selectedModel = useChatbotConfigStore(selectSelectedModel(configId));
75+
const selectedMcpServerIds = useChatbotConfigStore(selectSelectedMcpServerIds(configId));
76+
7677
const setSelectedModel = React.useCallback(
7778
(model: string) => {
7879
useChatbotConfigStore.getState().updateSelectedModel(configId, model);
@@ -94,10 +95,6 @@ const ChatbotPlayground: React.FC<ChatbotPlaygroundProps> = ({
9495
return statuses ? new Map(Object.entries(statuses)) : new Map();
9596
}, [location.state?.mcpServerStatuses]);
9697

97-
// Handle router state for MCP servers - initialize state with router value
98-
const [selectedMcpServerIds, setSelectedMcpServerIds] =
99-
React.useState<string[]>(mcpServersFromRoute);
100-
10198
// MCP hooks - fetch servers and manage statuses
10299
const {
103100
data: mcpServers = [],
@@ -119,13 +116,15 @@ const ChatbotPlayground: React.FC<ChatbotPlaygroundProps> = ({
119116
);
120117

121118
React.useEffect(() => {
122-
// Reset configuration
123-
useChatbotConfigStore.getState().resetConfiguration();
119+
// Reset configuration with initial values from router
120+
useChatbotConfigStore.getState().resetConfiguration({
121+
selectedMcpServerIds: mcpServersFromRoute,
122+
});
124123

125124
return () => {
126125
useChatbotConfigStore.getState().resetConfiguration();
127126
};
128-
}, [configId]);
127+
}, [configId, mcpServersFromRoute, selectedAAModel]);
129128

130129
// Clear router state after a brief delay to ensure child components have consumed it
131130
React.useEffect(() => {
@@ -183,6 +182,13 @@ const ChatbotPlayground: React.FC<ChatbotPlaygroundProps> = ({
183182
isFilesLoading: fileManagement.isLoading,
184183
});
185184

185+
// Get tool selections callback from store
186+
const getToolSelections = React.useCallback(
187+
(namespaceName: string, serverUrl: string) =>
188+
useChatbotConfigStore.getState().getToolSelections(configId, namespaceName, serverUrl),
189+
[configId],
190+
);
191+
186192
// TODO: This will need to be changed to an object or array when implementing compare mode
187193
const chatbotMessages = useChatbotMessages({
188194
modelId: selectedModel,
@@ -235,10 +241,7 @@ const ChatbotPlayground: React.FC<ChatbotPlaygroundProps> = ({
235241
alerts={{ uploadSuccessAlert, deleteSuccessAlert, errorAlert }}
236242
sourceManagement={sourceManagement}
237243
fileManagement={fileManagement}
238-
onMcpServersChange={setSelectedMcpServerIds}
239-
initialSelectedServerIds={mcpServersFromRoute}
240244
initialServerStatuses={mcpServerStatusesFromRoute}
241-
selectedServersCount={selectedMcpServerIds.length}
242245
mcpServers={mcpServers}
243246
mcpServersLoaded={mcpServersLoaded}
244247
mcpServersLoadError={mcpServersLoadError}

‎packages/gen-ai/frontend/src/app/Chatbot/components/ChatbotSettingsPanel.tsx‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import {
1919
selectStreamingEnabled,
2020
selectSelectedModel,
2121
selectGuardrailsEnabled,
22+
selectSelectedMcpServerIds,
2223
} from '~/app/Chatbot/store';
2324
import { UseSourceManagementReturn } from '~/app/Chatbot/hooks/useSourceManagement';
2425
import { UseFileManagementReturn } from '~/app/Chatbot/hooks/useFileManagement';
@@ -41,10 +42,7 @@ interface ChatbotSettingsPanelProps {
4142
};
4243
sourceManagement: UseSourceManagementReturn;
4344
fileManagement: UseFileManagementReturn;
44-
onMcpServersChange?: (serverIds: string[]) => void;
45-
initialSelectedServerIds?: string[];
4645
initialServerStatuses?: Map<string, ServerStatusInfo>;
47-
selectedServersCount: number;
4846
mcpServers: MCPServerFromAPI[];
4947
mcpServersLoaded: boolean;
5048
mcpServersLoadError?: Error | null;
@@ -61,10 +59,7 @@ const ChatbotSettingsPanel: React.FunctionComponent<ChatbotSettingsPanelProps> =
6159
alerts,
6260
sourceManagement,
6361
fileManagement,
64-
onMcpServersChange,
65-
initialSelectedServerIds,
6662
initialServerStatuses,
67-
selectedServersCount,
6863
mcpServers,
6964
mcpServersLoaded,
7065
mcpServersLoadError,
@@ -78,6 +73,7 @@ const ChatbotSettingsPanel: React.FunctionComponent<ChatbotSettingsPanelProps> =
7873
// Consume store directly using configId
7974
const systemInstruction = useChatbotConfigStore(selectSystemInstruction(configId));
8075
const temperature = useChatbotConfigStore(selectTemperature(configId));
76+
const selectedMcpServerIds = useChatbotConfigStore(selectSelectedMcpServerIds(configId));
8177
const isStreamingEnabled = useChatbotConfigStore(selectStreamingEnabled(configId));
8278
const selectedModel = useChatbotConfigStore(selectSelectedModel(configId));
8379
const guardrailsEnabled = useChatbotConfigStore(selectGuardrailsEnabled(configId));
@@ -223,9 +219,9 @@ const ChatbotSettingsPanel: React.FunctionComponent<ChatbotSettingsPanelProps> =
223219
<FlexItem>
224220
<TabTitleText>MCP</TabTitleText>
225221
</FlexItem>
226-
{selectedServersCount > 0 && (
222+
{selectedMcpServerIds.length > 0 && (
227223
<FlexItem>
228-
<Badge>{selectedServersCount}</Badge>
224+
<Badge>{selectedMcpServerIds.length}</Badge>
229225
</FlexItem>
230226
)}
231227
{showMcpToolsWarning && (
@@ -242,14 +238,13 @@ const ChatbotSettingsPanel: React.FunctionComponent<ChatbotSettingsPanelProps> =
242238
data-testid="chatbot-settings-page-tab-mcp"
243239
>
244240
<MCPTabContent
241+
configId={configId}
245242
mcpServers={mcpServers}
246243
mcpServersLoaded={mcpServersLoaded}
247244
mcpServersLoadError={mcpServersLoadError}
248245
mcpServerTokens={mcpServerTokens}
249246
onMcpServerTokensChange={onMcpServerTokensChange}
250247
checkMcpServerStatus={checkMcpServerStatus}
251-
onMcpServersChange={onMcpServersChange}
252-
initialSelectedServerIds={initialSelectedServerIds}
253248
initialServerStatuses={initialServerStatuses}
254249
activeToolsCount={activeToolsCount}
255250
onActiveToolsCountChange={setActiveToolsCount}

‎packages/gen-ai/frontend/src/app/Chatbot/components/settingsPanelTabs/MCPTabContent.tsx‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,29 +6,27 @@ import MCPServersPanel from '~/app/Chatbot/mcp/MCPServersPanel';
66
import TabContentWrapper from './TabContentWrapper';
77

88
interface MCPTabContentProps {
9+
configId: string;
910
mcpServers: MCPServerFromAPI[];
1011
mcpServersLoaded: boolean;
1112
mcpServersLoadError?: Error | null;
1213
mcpServerTokens: Map<string, TokenInfo>;
1314
onMcpServerTokensChange: (tokens: Map<string, TokenInfo>) => void;
1415
checkMcpServerStatus: (serverUrl: string, mcpBearerToken?: string) => Promise<ServerStatusInfo>;
15-
onMcpServersChange?: (serverIds: string[]) => void;
16-
initialSelectedServerIds?: string[];
1716
initialServerStatuses?: Map<string, ServerStatusInfo>;
1817
activeToolsCount: number;
1918
onActiveToolsCountChange: (count: number) => void;
2019
onToolsWarningChange: (show: boolean) => void;
2120
}
2221

2322
const MCPTabContent: React.FunctionComponent<MCPTabContentProps> = ({
23+
configId,
2424
mcpServers,
2525
mcpServersLoaded,
2626
mcpServersLoadError,
2727
mcpServerTokens,
2828
onMcpServerTokensChange,
2929
checkMcpServerStatus,
30-
onMcpServersChange,
31-
initialSelectedServerIds,
3230
initialServerStatuses,
3331
activeToolsCount,
3432
onActiveToolsCountChange,
@@ -48,14 +46,13 @@ const MCPTabContent: React.FunctionComponent<MCPTabContentProps> = ({
4846
titleTestId="mcp-servers-section-title"
4947
>
5048
<MCPServersPanel
49+
configId={configId}
5150
servers={mcpServers}
5251
serversLoaded={mcpServersLoaded}
5352
serversLoadError={mcpServersLoadError}
5453
serverTokens={mcpServerTokens}
5554
onServerTokensChange={onMcpServerTokensChange}
5655
checkServerStatus={checkMcpServerStatus}
57-
onSelectionChange={onMcpServersChange}
58-
initialSelectedServerIds={initialSelectedServerIds}
5956
initialServerStatuses={initialServerStatuses}
6057
onToolsWarningChange={onToolsWarningChange}
6158
onActiveToolsCountChange={onActiveToolsCountChange}

‎packages/gen-ai/frontend/src/app/Chatbot/components/settingsPanelTabs/__tests__/MCPTabContent.spec.tsx‎

Lines changed: 9 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -12,29 +12,21 @@ jest.mock('~/app/Chatbot/hooks/useDarkMode', () => ({
1212
jest.mock('~/app/Chatbot/mcp/MCPServersPanel', () => ({
1313
__esModule: true,
1414
default: ({
15+
configId,
1516
servers,
1617
serversLoaded,
1718
serversLoadError,
18-
onSelectionChange,
19-
initialSelectedServerIds,
2019
}: {
20+
configId: string;
2121
servers: MCPServerFromAPI[];
2222
serversLoaded: boolean;
2323
serversLoadError?: Error | null;
24-
onSelectionChange?: (serverIds: string[]) => void;
25-
initialSelectedServerIds?: string[];
2624
}) => (
2725
<div data-testid="mcp-servers-panel">
26+
<span data-testid="config-id">{configId}</span>
2827
<span data-testid="servers-count">{servers.length}</span>
2928
<span data-testid="servers-loaded">{serversLoaded.toString()}</span>
3029
<span data-testid="servers-error">{serversLoadError?.message || 'none'}</span>
31-
<span data-testid="initial-selected">{initialSelectedServerIds?.join(',') || 'none'}</span>
32-
<button
33-
data-testid="select-servers-button"
34-
onClick={() => onSelectionChange?.(['server-1', 'server-2'])}
35-
>
36-
Select Servers
37-
</button>
3830
</div>
3931
),
4032
}));
@@ -47,14 +39,13 @@ describe('MCPTabContent', () => {
4739
} as ServerStatusInfo);
4840

4941
const defaultProps = {
42+
configId: 'default',
5043
mcpServers: [] as MCPServerFromAPI[],
5144
mcpServersLoaded: true,
5245
mcpServersLoadError: null,
5346
mcpServerTokens: new Map<string, TokenInfo>(),
5447
onMcpServerTokensChange: jest.fn(),
5548
checkMcpServerStatus: mockCheckMcpServerStatus,
56-
onMcpServersChange: jest.fn(),
57-
initialSelectedServerIds: [] as string[],
5849
initialServerStatuses: new Map<string, ServerStatusInfo>(),
5950
activeToolsCount: 0,
6051
onActiveToolsCountChange: jest.fn(),
@@ -91,18 +82,12 @@ describe('MCPTabContent', () => {
9182
},
9283
];
9384

94-
render(
95-
<MCPTabContent
96-
{...defaultProps}
97-
mcpServers={servers}
98-
initialSelectedServerIds={['server-1']}
99-
/>,
100-
);
85+
render(<MCPTabContent {...defaultProps} mcpServers={servers} />);
10186

10287
expect(screen.getByTestId('mcp-servers-panel')).toBeInTheDocument();
88+
expect(screen.getByTestId('config-id')).toHaveTextContent('default');
10389
expect(screen.getByTestId('servers-count')).toHaveTextContent('2');
10490
expect(screen.getByTestId('servers-loaded')).toHaveTextContent('true');
105-
expect(screen.getByTestId('initial-selected')).toHaveTextContent('server-1');
10691
});
10792

10893
it('passes serversLoaded false when not loaded', () => {
@@ -151,30 +136,9 @@ describe('MCPTabContent', () => {
151136
expect(screen.getByTestId('active-tools-count-badge')).toHaveTextContent('42 tools enabled');
152137
});
153138

154-
it('calls onMcpServersChange when servers are selected', async () => {
155-
const mockOnMcpServersChange = jest.fn();
156-
render(<MCPTabContent {...defaultProps} onMcpServersChange={mockOnMcpServersChange} />);
139+
it('passes custom configId to MCPServersPanel', () => {
140+
render(<MCPTabContent {...defaultProps} configId="test-config" />);
157141

158-
const selectButton = screen.getByTestId('select-servers-button');
159-
selectButton.click();
160-
161-
expect(mockOnMcpServersChange).toHaveBeenCalledWith(['server-1', 'server-2']);
162-
});
163-
164-
it('renders with empty initial selected server ids', () => {
165-
render(<MCPTabContent {...defaultProps} initialSelectedServerIds={[]} />);
166-
167-
expect(screen.getByTestId('initial-selected')).toHaveTextContent('none');
168-
});
169-
170-
it('renders with multiple initial selected server ids', () => {
171-
render(
172-
<MCPTabContent
173-
{...defaultProps}
174-
initialSelectedServerIds={['server-a', 'server-b', 'server-c']}
175-
/>,
176-
);
177-
178-
expect(screen.getByTestId('initial-selected')).toHaveTextContent('server-a,server-b,server-c');
142+
expect(screen.getByTestId('config-id')).toHaveTextContent('test-config');
179143
});
180144
});

‎packages/gen-ai/frontend/src/app/Chatbot/mcp/MCPServerToolsModal.tsx‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,14 @@ import {
2020
} from '@odh-dashboard/internal/concepts/analyticsTracking/segmentIOUtils';
2121
import { TrackingOutcome } from '@odh-dashboard/internal/concepts/analyticsTracking/trackingProperties';
2222
import { useMCPServerTools } from '~/app/hooks/useMCPServerTools';
23-
import { useMCPToolSelections } from '~/app/hooks/useMCPToolSelections';
2423
import { GenAiContext } from '~/app/context/GenAiContext';
2524
import { MCPServer, MCPTool } from '~/app/types';
25+
import { useChatbotConfigStore } from '~/app/Chatbot/store';
2626
import MCPToolsColumns from './MCPToolsColumns';
2727
import MCPToolsTableRow from './MCPToolsTableRow';
2828

2929
interface MCPServerToolsModalProps {
30+
configId: string;
3031
isOpen: boolean;
3132
onClose: () => void;
3233
server: MCPServer;
@@ -36,13 +37,28 @@ interface MCPServerToolsModalProps {
3637
const MCP_TOOLS_EVENT_NAME = 'Playground MCP Tools';
3738

3839
const MCPServerToolsModal: React.FC<MCPServerToolsModalProps> = ({
40+
configId,
3941
isOpen,
4042
onClose,
4143
server,
4244
mcpBearerToken,
4345
}) => {
4446
const { namespace } = React.useContext(GenAiContext);
45-
const { getToolSelections, saveToolSelections } = useMCPToolSelections();
47+
48+
// Get tool selections functions from store
49+
const getToolSelections = React.useCallback(
50+
(namespaceName: string, serverUrl: string) =>
51+
useChatbotConfigStore.getState().getToolSelections(configId, namespaceName, serverUrl),
52+
[configId],
53+
);
54+
const saveToolSelections = React.useCallback(
55+
(namespaceName: string, serverUrl: string, toolNames: string[] | undefined) => {
56+
useChatbotConfigStore
57+
.getState()
58+
.saveToolSelections(configId, namespaceName, serverUrl, toolNames);
59+
},
60+
[configId],
61+
);
4662

4763
const {
4864
tools: apiTools,

0 commit comments

Comments
 (0)