🗂️ refactor: Make MCPSubMenu consistent with MCPSelect (#8650)

- Refactored MCPSelect and MCPSubMenu components to utilize a new custom hook, `useMCPServerManager`, for improved state management and server initialization logic.
- Added functionality to handle simultaneous MCP server initialization requests, including cancellation and user notifications.
- Updated translation files to include new messages for initialization cancellation.
- Improved the configuration dialog handling for MCP servers, streamlining the user experience when managing server settings.
This commit is contained in:
Dustin Healy 2025-07-25 11:51:42 -07:00 committed by GitHub
parent cd436dc6a8
commit 545a909953
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 503 additions and 424 deletions

View file

@ -331,7 +331,8 @@ router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => {
logger.info(`[MCP Reinitialize] Reinitializing server: ${serverName}`);
const config = await loadCustomConfig();
const printConfig = false;
const config = await loadCustomConfig(printConfig);
if (!config || !config.mcpServers || !config.mcpServers[serverName]) {
return res.status(404).json({
error: `MCP server '${serverName}' not found in configuration`,

View file

@ -1,114 +1,21 @@
import { useQueryClient } from '@tanstack/react-query';
import { Constants, QueryKeys } from 'librechat-data-provider';
import type { TUpdateUserPlugins, TPlugin } from 'librechat-data-provider';
import React, { memo, useCallback, useState, useMemo, useRef } from 'react';
import { useUpdateUserPluginsMutation } from 'librechat-data-provider/react-query';
import MCPConfigDialog, { ConfigFieldDetail } from '~/components/ui/MCP/MCPConfigDialog';
import { useMCPServerInitialization } from '~/hooks/MCP/useMCPServerInitialization';
import React, { memo, useCallback } from 'react';
import MCPConfigDialog from '~/components/ui/MCP/MCPConfigDialog';
import MCPServerStatusIcon from '~/components/ui/MCP/MCPServerStatusIcon';
import { useToastContext, useBadgeRowContext } from '~/Providers';
import MultiSelect from '~/components/ui/MultiSelect';
import { MCPIcon } from '~/components/svg';
import { useLocalize } from '~/hooks';
import { useMCPServerManager } from '~/hooks/MCP/useMCPServerManager';
function MCPSelect() {
const localize = useLocalize();
const { showToast } = useToastContext();
const { mcpSelect, startupConfig } = useBadgeRowContext();
const { mcpValues, setMCPValues, mcpToolDetails, isPinned } = mcpSelect;
// Get all configured MCP servers from config that allow chat menu
const configuredServers = useMemo(() => {
if (!startupConfig?.mcpServers) {
return [];
}
return Object.entries(startupConfig.mcpServers)
.filter(([, config]) => config.chatMenu !== false)
.map(([serverName]) => serverName);
}, [startupConfig?.mcpServers]);
const [isConfigModalOpen, setIsConfigModalOpen] = useState(false);
const [selectedToolForConfig, setSelectedToolForConfig] = useState<TPlugin | null>(null);
const previousFocusRef = useRef<HTMLElement | null>(null);
const queryClient = useQueryClient();
const updateUserPluginsMutation = useUpdateUserPluginsMutation({
onSuccess: async () => {
showToast({ message: localize('com_nav_mcp_vars_updated'), status: 'success' });
// tools so we dont leave tools available for use in chat if we revoke and thus kill mcp server
// auth values so customUserVars flags are updated in customUserVarsSection
// connection status so connection indicators are updated in the dropdown
await Promise.all([
queryClient.refetchQueries([QueryKeys.tools]),
queryClient.refetchQueries([QueryKeys.mcpAuthValues]),
queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]),
]);
},
onError: (error: unknown) => {
console.error('Error updating MCP auth:', error);
showToast({
message: localize('com_nav_mcp_vars_update_error'),
status: 'error',
});
},
});
// Use the shared initialization hook
const { initializeServer, isInitializing, connectionStatus, cancelOAuthFlow, isCancellable } =
useMCPServerInitialization({
onSuccess: (serverName) => {
// Add to selected values after successful initialization
const currentValues = mcpValues ?? [];
if (!currentValues.includes(serverName)) {
setMCPValues([...currentValues, serverName]);
}
},
onError: (serverName) => {
// Find the tool/server configuration
const tool = mcpToolDetails?.find((t) => t.name === serverName);
const serverConfig = startupConfig?.mcpServers?.[serverName];
const serverStatus = connectionStatus[serverName];
// Check if this server would show a config button
const hasAuthConfig =
(tool?.authConfig && tool.authConfig.length > 0) ||
(serverConfig?.customUserVars && Object.keys(serverConfig.customUserVars).length > 0);
// Only open dialog if the server would have shown a config button
// (disconnected/error states always show button, connected only shows if hasAuthConfig)
const wouldShowButton =
!serverStatus ||
serverStatus.connectionState === 'disconnected' ||
serverStatus.connectionState === 'error' ||
(serverStatus.connectionState === 'connected' && hasAuthConfig);
if (!wouldShowButton) {
return; // Don't open dialog if no button would be shown
}
// Create tool object if it doesn't exist
const configTool = tool || {
name: serverName,
pluginKey: `${Constants.mcp_prefix}${serverName}`,
authConfig: serverConfig?.customUserVars
? Object.entries(serverConfig.customUserVars).map(([key, config]) => ({
authField: key,
label: config.title,
description: config.description,
}))
: [],
authenticated: false,
};
previousFocusRef.current = document.activeElement as HTMLElement;
// Open the config dialog on error
setSelectedToolForConfig(configTool);
setIsConfigModalOpen(true);
},
});
const {
configuredServers,
mcpValues,
isPinned,
placeholderText,
batchToggleServers,
getServerStatusIconProps,
getConfigDialogProps,
localize,
} = useMCPServerManager();
const renderSelectedValues = useCallback(
(values: string[], placeholder?: string) => {
@ -123,137 +30,9 @@ function MCPSelect() {
[localize],
);
const handleConfigSave = useCallback(
(targetName: string, authData: Record<string, string>) => {
if (selectedToolForConfig && selectedToolForConfig.name === targetName) {
// Use the pluginKey directly since it's already in the correct format
console.log(
`[MCP Select] Saving config for ${targetName}, pluginKey: ${`${Constants.mcp_prefix}${targetName}`}`,
);
const payload: TUpdateUserPlugins = {
pluginKey: `${Constants.mcp_prefix}${targetName}`,
action: 'install',
auth: authData,
};
updateUserPluginsMutation.mutate(payload);
}
},
[selectedToolForConfig, updateUserPluginsMutation],
);
const handleConfigRevoke = useCallback(
(targetName: string) => {
if (selectedToolForConfig && selectedToolForConfig.name === targetName) {
// Use the pluginKey directly since it's already in the correct format
const payload: TUpdateUserPlugins = {
pluginKey: `${Constants.mcp_prefix}${targetName}`,
action: 'uninstall',
auth: {},
};
updateUserPluginsMutation.mutate(payload);
// Remove the server from selected values after revoke
const currentValues = mcpValues ?? [];
const filteredValues = currentValues.filter((name) => name !== targetName);
setMCPValues(filteredValues);
}
},
[selectedToolForConfig, updateUserPluginsMutation, mcpValues, setMCPValues],
);
const handleSave = useCallback(
(authData: Record<string, string>) => {
if (selectedToolForConfig) {
handleConfigSave(selectedToolForConfig.name, authData);
}
},
[selectedToolForConfig, handleConfigSave],
);
const handleRevoke = useCallback(() => {
if (selectedToolForConfig) {
handleConfigRevoke(selectedToolForConfig.name);
}
}, [selectedToolForConfig, handleConfigRevoke]);
const handleDialogOpenChange = useCallback((open: boolean) => {
setIsConfigModalOpen(open);
// Restore focus when dialog closes
if (!open && previousFocusRef.current) {
// Use setTimeout to ensure the dialog has fully closed before restoring focus
setTimeout(() => {
if (previousFocusRef.current && typeof previousFocusRef.current.focus === 'function') {
previousFocusRef.current.focus();
}
previousFocusRef.current = null;
}, 0);
}
}, []);
// Get connection status for all MCP servers (now from hook)
// Remove the duplicate useMCPConnectionStatusQuery since it's in the hook
// Modified setValue function that attempts to initialize disconnected servers
const filteredSetMCPValues = useCallback(
(values: string[]) => {
// Separate connected and disconnected servers
const connectedServers: string[] = [];
const disconnectedServers: string[] = [];
values.forEach((serverName) => {
const serverStatus = connectionStatus[serverName];
if (serverStatus?.connectionState === 'connected') {
connectedServers.push(serverName);
} else {
disconnectedServers.push(serverName);
}
});
// Only set connected servers as selected values
setMCPValues(connectedServers);
// Attempt to initialize each disconnected server (once)
disconnectedServers.forEach((serverName) => {
initializeServer(serverName);
});
},
[connectionStatus, setMCPValues, initializeServer],
);
const renderItemContent = useCallback(
(serverName: string, defaultContent: React.ReactNode) => {
const tool = mcpToolDetails?.find((t) => t.name === serverName);
const serverStatus = connectionStatus[serverName];
const serverConfig = startupConfig?.mcpServers?.[serverName];
const handleConfigClick = (e: React.MouseEvent) => {
e.stopPropagation();
e.preventDefault();
previousFocusRef.current = document.activeElement as HTMLElement;
const configTool = tool || {
name: serverName,
pluginKey: `${Constants.mcp_prefix}${serverName}`,
authConfig: serverConfig?.customUserVars
? Object.entries(serverConfig.customUserVars).map(([key, config]) => ({
authField: key,
label: config.title,
description: config.description,
}))
: [],
authenticated: false,
};
setSelectedToolForConfig(configTool);
setIsConfigModalOpen(true);
};
const handleCancelClick = (e: React.MouseEvent) => {
e.stopPropagation();
e.preventDefault();
cancelOAuthFlow(serverName);
};
const statusIconProps = getServerStatusIconProps(serverName);
// Common wrapper for the main content (check mark + text)
// Ensures Check & Text are adjacent and the group takes available space.
@ -267,22 +46,7 @@ function MCPSelect() {
</button>
);
// Check if this server has customUserVars to configure
const hasCustomUserVars =
serverConfig?.customUserVars && Object.keys(serverConfig.customUserVars).length > 0;
const statusIcon = (
<MCPServerStatusIcon
serverName={serverName}
serverStatus={serverStatus}
tool={tool}
onConfigClick={handleConfigClick}
isInitializing={isInitializing(serverName)}
canCancel={isCancellable(serverName)}
onCancel={handleCancelClick}
hasCustomUserVars={hasCustomUserVars}
/>
);
const statusIcon = statusIconProps && <MCPServerStatusIcon {...statusIconProps} />;
if (statusIcon) {
return (
@ -295,14 +59,7 @@ function MCPSelect() {
return mainContentWrapper;
},
[
isInitializing,
isCancellable,
mcpToolDetails,
cancelOAuthFlow,
connectionStatus,
startupConfig?.mcpServers,
],
[getServerStatusIconProps],
);
// Don't render if no servers are selected and not pinned
@ -315,14 +72,14 @@ function MCPSelect() {
return null;
}
const placeholderText =
startupConfig?.interface?.mcpServers?.placeholder || localize('com_ui_mcp_servers');
const configDialogProps = getConfigDialogProps();
return (
<>
<MultiSelect
items={configuredServers}
selectedValues={mcpValues ?? []}
setSelectedValues={filteredSetMCPValues}
setSelectedValues={batchToggleServers}
defaultSelectedValues={mcpValues ?? []}
renderSelectedValues={renderSelectedValues}
renderItemContent={renderItemContent}
@ -333,39 +90,7 @@ function MCPSelect() {
selectItemsClassName="border border-blue-600/50 bg-blue-500/10 hover:bg-blue-700/10"
selectClassName="group relative inline-flex items-center justify-center md:justify-start gap-1.5 rounded-full border border-border-medium text-sm font-medium transition-all md:w-full size-9 p-2 md:p-3 bg-transparent shadow-sm hover:bg-surface-hover hover:shadow-md active:shadow-inner"
/>
{selectedToolForConfig && (
<MCPConfigDialog
serverName={selectedToolForConfig.name}
serverStatus={connectionStatus[selectedToolForConfig.name]}
isOpen={isConfigModalOpen}
onOpenChange={handleDialogOpenChange}
fieldsSchema={(() => {
const schema: Record<string, ConfigFieldDetail> = {};
if (selectedToolForConfig?.authConfig) {
selectedToolForConfig.authConfig.forEach((field) => {
schema[field.authField] = {
title: field.label,
description: field.description,
};
});
}
return schema;
})()}
initialValues={(() => {
const initial: Record<string, string> = {};
// Note: Actual initial values might need to be fetched if they are stored user-specifically
if (selectedToolForConfig?.authConfig) {
selectedToolForConfig.authConfig.forEach((field) => {
initial[field.authField] = ''; // Or fetched value
});
}
return initial;
})()}
onSave={handleSave}
onRevoke={handleRevoke}
isSubmitting={updateUserPluginsMutation.isLoading}
/>
)}
{configDialogProps && <MCPConfigDialog {...configDialogProps} />}
</>
);
}

View file

@ -2,28 +2,26 @@ import React from 'react';
import * as Ariakit from '@ariakit/react';
import { ChevronRight } from 'lucide-react';
import { PinIcon, MCPIcon } from '~/components/svg';
import { useLocalize } from '~/hooks';
import MCPConfigDialog from '~/components/ui/MCP/MCPConfigDialog';
import MCPServerStatusIcon from '~/components/ui/MCP/MCPServerStatusIcon';
import { useMCPServerManager } from '~/hooks/MCP/useMCPServerManager';
import { cn } from '~/utils';
interface MCPSubMenuProps {
isMCPPinned: boolean;
setIsMCPPinned: (value: boolean) => void;
mcpValues?: string[];
mcpServerNames: string[];
handleMCPToggle: (serverName: string) => void;
placeholder?: string;
}
const MCPSubMenu = ({
mcpValues,
isMCPPinned,
mcpServerNames,
setIsMCPPinned,
handleMCPToggle,
placeholder,
...props
}: MCPSubMenuProps) => {
const localize = useLocalize();
const MCPSubMenu = ({ placeholder, ...props }: MCPSubMenuProps) => {
const {
configuredServers,
mcpValues,
isPinned,
setIsPinned,
placeholderText,
toggleServerSelection,
getServerStatusIconProps,
getConfigDialogProps,
} = useMCPServerManager();
const menuStore = Ariakit.useMenuStore({
focusLoop: true,
@ -31,72 +29,96 @@ const MCPSubMenu = ({
placement: 'right',
});
// Don't render if no MCP servers are configured
if (!configuredServers || configuredServers.length === 0) {
return null;
}
const configDialogProps = getConfigDialogProps();
return (
<Ariakit.MenuProvider store={menuStore}>
<Ariakit.MenuItem
{...props}
render={
<Ariakit.MenuButton
onClick={(e: React.MouseEvent<HTMLButtonElement>) => {
e.stopPropagation();
menuStore.toggle();
}}
className="flex w-full cursor-pointer items-center justify-between rounded-lg p-2 hover:bg-surface-hover"
/>
}
>
<div className="flex items-center gap-2">
<MCPIcon className="icon-md" />
<span>{placeholder || localize('com_ui_mcp_servers')}</span>
<ChevronRight className="ml-auto h-3 w-3" />
</div>
<button
type="button"
onClick={(e) => {
e.stopPropagation();
setIsMCPPinned(!isMCPPinned);
}}
className={cn(
'rounded p-1 transition-all duration-200',
'hover:bg-surface-tertiary hover:shadow-sm',
!isMCPPinned && 'text-text-secondary hover:text-text-primary',
)}
aria-label={isMCPPinned ? 'Unpin' : 'Pin'}
<>
<Ariakit.MenuProvider store={menuStore}>
<Ariakit.MenuItem
{...props}
render={
<Ariakit.MenuButton
onClick={(e: React.MouseEvent<HTMLButtonElement>) => {
e.stopPropagation();
menuStore.toggle();
}}
className="flex w-full cursor-pointer items-center justify-between rounded-lg p-2 hover:bg-surface-hover"
/>
}
>
<div className="h-4 w-4">
<PinIcon unpin={isMCPPinned} />
<div className="flex items-center gap-2">
<MCPIcon className="icon-md" />
<span>{placeholder || placeholderText}</span>
<ChevronRight className="ml-auto h-3 w-3" />
</div>
</button>
</Ariakit.MenuItem>
<Ariakit.Menu
portal={true}
unmountOnHide={true}
className={cn(
'animate-popover-left z-50 ml-3 flex min-w-[200px] flex-col rounded-xl',
'border border-border-light bg-surface-secondary p-1 shadow-lg',
)}
>
{mcpServerNames.map((serverName) => (
<Ariakit.MenuItem
key={serverName}
onClick={(event) => {
event.preventDefault();
handleMCPToggle(serverName);
<button
type="button"
onClick={(e) => {
e.stopPropagation();
setIsPinned(!isPinned);
}}
className={cn(
'flex items-center gap-2 rounded-lg px-2 py-1.5 text-text-primary hover:cursor-pointer',
'scroll-m-1 outline-none transition-colors',
'hover:bg-black/[0.075] dark:hover:bg-white/10',
'data-[active-item]:bg-black/[0.075] dark:data-[active-item]:bg-white/10',
'w-full min-w-0 text-sm',
'rounded p-1 transition-all duration-200',
'hover:bg-surface-tertiary hover:shadow-sm',
!isPinned && 'text-text-secondary hover:text-text-primary',
)}
aria-label={isPinned ? 'Unpin' : 'Pin'}
>
<Ariakit.MenuItemCheck checked={mcpValues?.includes(serverName) ?? false} />
<span>{serverName}</span>
</Ariakit.MenuItem>
))}
</Ariakit.Menu>
</Ariakit.MenuProvider>
<div className="h-4 w-4">
<PinIcon unpin={isPinned} />
</div>
</button>
</Ariakit.MenuItem>
<Ariakit.Menu
portal={true}
unmountOnHide={true}
className={cn(
'animate-popover-left z-50 ml-3 flex min-w-[200px] flex-col rounded-xl',
'border border-border-light bg-surface-secondary p-1 shadow-lg',
)}
>
{configuredServers.map((serverName) => {
const statusIconProps = getServerStatusIconProps(serverName);
const isSelected = mcpValues?.includes(serverName) ?? false;
const statusIcon = statusIconProps && <MCPServerStatusIcon {...statusIconProps} />;
return (
<Ariakit.MenuItem
key={serverName}
onClick={(event) => {
event.preventDefault();
toggleServerSelection(serverName);
}}
className={cn(
'flex items-center gap-2 rounded-lg px-2 py-1.5 text-text-primary hover:cursor-pointer',
'scroll-m-1 outline-none transition-colors',
'hover:bg-black/[0.075] dark:hover:bg-white/10',
'data-[active-item]:bg-black/[0.075] dark:data-[active-item]:bg-white/10',
'w-full min-w-0 justify-between text-sm',
)}
>
<button
type="button"
className="flex flex-grow items-center gap-2 rounded bg-transparent p-0 text-left transition-colors focus:outline-none"
tabIndex={0}
>
<Ariakit.MenuItemCheck checked={isSelected} />
<span>{serverName}</span>
</button>
{statusIcon && <div className="ml-2 flex items-center">{statusIcon}</div>}
</Ariakit.MenuItem>
);
})}
</Ariakit.Menu>
</Ariakit.MenuProvider>
{configDialogProps && <MCPConfigDialog {...configDialogProps} />}
</>
);
};

View file

@ -55,12 +55,7 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => {
} = codeInterpreter;
const { isPinned: isFileSearchPinned, setIsPinned: setIsFileSearchPinned } = fileSearch;
const { isPinned: isArtifactsPinned, setIsPinned: setIsArtifactsPinned } = artifacts;
const {
mcpValues,
mcpServerNames,
isPinned: isMCPPinned,
setIsPinned: setIsMCPPinned,
} = mcpSelect;
const { mcpServerNames } = mcpSelect;
const canUseWebSearch = useHasAccess({
permissionType: PermissionTypes.WEB_SEARCH,
@ -130,17 +125,6 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => {
}
}, [artifacts]);
const handleMCPToggle = useCallback(
(serverName: string) => {
const currentValues = mcpSelect.mcpValues ?? [];
const newValues = currentValues.includes(serverName)
? currentValues.filter((v) => v !== serverName)
: [...currentValues, serverName];
mcpSelect.setMCPValues(newValues);
},
[mcpSelect],
);
const mcpPlaceholder = startupConfig?.interface?.mcpServers?.placeholder;
const dropdownItems: MenuItemProps[] = [];
@ -305,17 +289,7 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => {
if (mcpServerNames && mcpServerNames.length > 0) {
dropdownItems.push({
hideOnClick: false,
render: (props) => (
<MCPSubMenu
{...props}
mcpValues={mcpValues}
isMCPPinned={isMCPPinned}
placeholder={mcpPlaceholder}
mcpServerNames={mcpServerNames}
setIsMCPPinned={setIsMCPPinned}
handleMCPToggle={handleMCPToggle}
/>
),
render: (props) => <MCPSubMenu {...props} placeholder={mcpPlaceholder} />,
});
}

View file

@ -37,6 +37,9 @@ export function useMCPServerInitialization(options?: UseMCPServerInitializationO
// Main initialization mutation
const reinitializeMutation = useReinitializeMCPServerMutation();
// Track which server is currently being processed
const [currentProcessingServer, setCurrentProcessingServer] = useState<string | null>(null);
// Cancel OAuth mutation
const cancelOAuthMutation = useCancelMCPOAuthMutation();
@ -184,12 +187,32 @@ export function useMCPServerInitialization(options?: UseMCPServerInitializationO
return;
}
if (connectionStatus[serverName]?.requiresOAuth) {
setCancellableServers((prev) => new Set(prev).add(serverName));
}
// Add to initializing set
setInitializingServers((prev) => new Set(prev).add(serverName));
// Trigger initialization
// If there's already a server being processed, that one will be cancelled
if (currentProcessingServer && currentProcessingServer !== serverName) {
// Clean up the cancelled server's state immediately
showToast({
message: localize('com_ui_mcp_init_cancelled', { 0: currentProcessingServer }),
status: 'warning',
});
cleanupOAuthState(currentProcessingServer);
}
// Track the current server being processed
setCurrentProcessingServer(serverName);
reinitializeMutation.mutate(serverName, {
onSuccess: (response: any) => {
// Clear current processing server
setCurrentProcessingServer(null);
if (response.success) {
if (response.oauthRequired && response.oauthUrl) {
// OAuth required - store URL and start polling
@ -238,40 +261,45 @@ export function useMCPServerInitialization(options?: UseMCPServerInitializationO
}
},
onError: (error: any) => {
console.error('Error initializing MCP server:', error);
showToast({
message: localize('com_ui_mcp_init_failed'),
status: 'error',
});
// Remove from initializing on error
setInitializingServers((prev) => {
const newSet = new Set(prev);
newSet.delete(serverName);
return newSet;
});
// Remove from OAuth tracking
setOauthPollingServers((prev) => {
const newMap = new Map(prev);
newMap.delete(serverName);
return newMap;
});
setOauthStartTimes((prev) => {
const newMap = new Map(prev);
newMap.delete(serverName);
return newMap;
});
console.error(`Error initializing MCP server ${serverName}:`, error);
setCurrentProcessingServer(null);
const isCancelled =
error?.name === 'CanceledError' ||
error?.code === 'ERR_CANCELED' ||
error?.message?.includes('cancel') ||
error?.message?.includes('abort');
if (isCancelled) {
showToast({
message: localize('com_ui_mcp_init_cancelled', { 0: serverName }),
status: 'warning',
});
} else {
showToast({
message: localize('com_ui_mcp_init_failed'),
status: 'error',
});
}
// Clean up OAuth state using helper function
cleanupOAuthState(serverName);
// Call optional error callback
options?.onError?.(serverName, error);
},
});
},
[
initializingServers,
connectionStatus,
currentProcessingServer,
reinitializeMutation,
showToast,
localize,
handleSuccessfulConnection,
initializingServers,
cleanupOAuthState,
options,
handleSuccessfulConnection,
],
);

View file

@ -0,0 +1,328 @@
import { useQueryClient } from '@tanstack/react-query';
import { Constants, QueryKeys } from 'librechat-data-provider';
import { useCallback, useState, useMemo, useRef } from 'react';
import { useUpdateUserPluginsMutation } from 'librechat-data-provider/react-query';
import { useMCPServerInitialization } from '~/hooks/MCP/useMCPServerInitialization';
import type { ConfigFieldDetail } from '~/components/ui/MCP/MCPConfigDialog';
import type { TUpdateUserPlugins, TPlugin } from 'librechat-data-provider';
import { useToastContext, useBadgeRowContext } from '~/Providers';
import { useLocalize } from '~/hooks';
export function useMCPServerManager() {
const localize = useLocalize();
const { showToast } = useToastContext();
const { mcpSelect, startupConfig } = useBadgeRowContext();
const { mcpValues, setMCPValues, mcpToolDetails, isPinned, setIsPinned } = mcpSelect;
const configuredServers = useMemo(() => {
if (!startupConfig?.mcpServers) {
return [];
}
return Object.entries(startupConfig.mcpServers)
.filter(([, config]) => config.chatMenu !== false)
.map(([serverName]) => serverName);
}, [startupConfig?.mcpServers]);
const [isConfigModalOpen, setIsConfigModalOpen] = useState(false);
const [selectedToolForConfig, setSelectedToolForConfig] = useState<TPlugin | null>(null);
const previousFocusRef = useRef<HTMLElement | null>(null);
const queryClient = useQueryClient();
const updateUserPluginsMutation = useUpdateUserPluginsMutation({
onSuccess: async () => {
showToast({ message: localize('com_nav_mcp_vars_updated'), status: 'success' });
await Promise.all([
queryClient.refetchQueries([QueryKeys.tools]),
queryClient.refetchQueries([QueryKeys.mcpAuthValues]),
queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]),
]);
},
onError: (error: unknown) => {
console.error('Error updating MCP auth:', error);
showToast({
message: localize('com_nav_mcp_vars_update_error'),
status: 'error',
});
},
});
const { initializeServer, isInitializing, connectionStatus, cancelOAuthFlow, isCancellable } =
useMCPServerInitialization({
onSuccess: (serverName) => {
const currentValues = mcpValues ?? [];
if (!currentValues.includes(serverName)) {
setMCPValues([...currentValues, serverName]);
}
},
onError: (serverName) => {
const tool = mcpToolDetails?.find((t) => t.name === serverName);
const serverConfig = startupConfig?.mcpServers?.[serverName];
const serverStatus = connectionStatus[serverName];
const hasAuthConfig =
(tool?.authConfig && tool.authConfig.length > 0) ||
(serverConfig?.customUserVars && Object.keys(serverConfig.customUserVars).length > 0);
const wouldShowButton =
!serverStatus ||
serverStatus.connectionState === 'disconnected' ||
serverStatus.connectionState === 'error' ||
(serverStatus.connectionState === 'connected' && hasAuthConfig);
if (!wouldShowButton) {
return;
}
const configTool = tool || {
name: serverName,
pluginKey: `${Constants.mcp_prefix}${serverName}`,
authConfig: serverConfig?.customUserVars
? Object.entries(serverConfig.customUserVars).map(([key, config]) => ({
authField: key,
label: config.title,
description: config.description,
}))
: [],
authenticated: false,
};
previousFocusRef.current = document.activeElement as HTMLElement;
setSelectedToolForConfig(configTool);
setIsConfigModalOpen(true);
},
});
const handleConfigSave = useCallback(
(targetName: string, authData: Record<string, string>) => {
if (selectedToolForConfig && selectedToolForConfig.name === targetName) {
const payload: TUpdateUserPlugins = {
pluginKey: `${Constants.mcp_prefix}${targetName}`,
action: 'install',
auth: authData,
};
updateUserPluginsMutation.mutate(payload);
}
},
[selectedToolForConfig, updateUserPluginsMutation],
);
const handleConfigRevoke = useCallback(
(targetName: string) => {
if (selectedToolForConfig && selectedToolForConfig.name === targetName) {
const payload: TUpdateUserPlugins = {
pluginKey: `${Constants.mcp_prefix}${targetName}`,
action: 'uninstall',
auth: {},
};
updateUserPluginsMutation.mutate(payload);
const currentValues = mcpValues ?? [];
const filteredValues = currentValues.filter((name) => name !== targetName);
setMCPValues(filteredValues);
}
},
[selectedToolForConfig, updateUserPluginsMutation, mcpValues, setMCPValues],
);
const handleSave = useCallback(
(authData: Record<string, string>) => {
if (selectedToolForConfig) {
handleConfigSave(selectedToolForConfig.name, authData);
}
},
[selectedToolForConfig, handleConfigSave],
);
const handleRevoke = useCallback(() => {
if (selectedToolForConfig) {
handleConfigRevoke(selectedToolForConfig.name);
}
}, [selectedToolForConfig, handleConfigRevoke]);
const handleDialogOpenChange = useCallback((open: boolean) => {
setIsConfigModalOpen(open);
if (!open && previousFocusRef.current) {
setTimeout(() => {
if (previousFocusRef.current && typeof previousFocusRef.current.focus === 'function') {
previousFocusRef.current.focus();
}
previousFocusRef.current = null;
}, 0);
}
}, []);
const toggleServerSelection = useCallback(
(serverName: string) => {
const currentValues = mcpValues ?? [];
const serverStatus = connectionStatus[serverName];
if (currentValues.includes(serverName)) {
const filteredValues = currentValues.filter((name) => name !== serverName);
setMCPValues(filteredValues);
} else {
if (serverStatus?.connectionState === 'connected') {
setMCPValues([...currentValues, serverName]);
} else {
initializeServer(serverName);
}
}
},
[connectionStatus, mcpValues, setMCPValues, initializeServer],
);
const batchToggleServers = useCallback(
(serverNames: string[]) => {
const connectedServers: string[] = [];
const disconnectedServers: string[] = [];
serverNames.forEach((serverName) => {
const serverStatus = connectionStatus[serverName];
if (serverStatus?.connectionState === 'connected') {
connectedServers.push(serverName);
} else {
disconnectedServers.push(serverName);
}
});
setMCPValues(connectedServers);
disconnectedServers.forEach((serverName) => {
initializeServer(serverName);
});
},
[connectionStatus, setMCPValues, initializeServer],
);
const getServerStatusIconProps = useCallback(
(serverName: string) => {
const tool = mcpToolDetails?.find((t) => t.name === serverName);
const serverStatus = connectionStatus[serverName];
const serverConfig = startupConfig?.mcpServers?.[serverName];
const handleConfigClick = (e: React.MouseEvent) => {
e.stopPropagation();
e.preventDefault();
previousFocusRef.current = document.activeElement as HTMLElement;
const configTool = tool || {
name: serverName,
pluginKey: `${Constants.mcp_prefix}${serverName}`,
authConfig: serverConfig?.customUserVars
? Object.entries(serverConfig.customUserVars).map(([key, config]) => ({
authField: key,
label: config.title,
description: config.description,
}))
: [],
authenticated: false,
};
setSelectedToolForConfig(configTool);
setIsConfigModalOpen(true);
};
const handleCancelClick = (e: React.MouseEvent) => {
e.stopPropagation();
e.preventDefault();
cancelOAuthFlow(serverName);
};
const hasCustomUserVars =
serverConfig?.customUserVars && Object.keys(serverConfig.customUserVars).length > 0;
return {
serverName,
serverStatus,
tool,
onConfigClick: handleConfigClick,
isInitializing: isInitializing(serverName),
canCancel: isCancellable(serverName),
onCancel: handleCancelClick,
hasCustomUserVars,
};
},
[
mcpToolDetails,
connectionStatus,
startupConfig?.mcpServers,
isInitializing,
isCancellable,
cancelOAuthFlow,
],
);
const placeholderText = useMemo(
() => startupConfig?.interface?.mcpServers?.placeholder || localize('com_ui_mcp_servers'),
[startupConfig?.interface?.mcpServers?.placeholder, localize],
);
const getConfigDialogProps = useCallback(() => {
if (!selectedToolForConfig) return null;
const fieldsSchema: Record<string, ConfigFieldDetail> = {};
if (selectedToolForConfig?.authConfig) {
selectedToolForConfig.authConfig.forEach((field) => {
fieldsSchema[field.authField] = {
title: field.label || field.authField,
description: field.description,
};
});
}
const initialValues: Record<string, string> = {};
if (selectedToolForConfig?.authConfig) {
selectedToolForConfig.authConfig.forEach((field) => {
initialValues[field.authField] = '';
});
}
return {
serverName: selectedToolForConfig.name,
serverStatus: connectionStatus[selectedToolForConfig.name],
isOpen: isConfigModalOpen,
onOpenChange: handleDialogOpenChange,
fieldsSchema,
initialValues,
onSave: handleSave,
onRevoke: handleRevoke,
isSubmitting: updateUserPluginsMutation.isLoading,
};
}, [
selectedToolForConfig,
connectionStatus,
isConfigModalOpen,
handleDialogOpenChange,
handleSave,
handleRevoke,
updateUserPluginsMutation.isLoading,
]);
return {
// Data
configuredServers,
mcpValues,
mcpToolDetails,
isPinned,
setIsPinned,
startupConfig,
connectionStatus,
placeholderText,
// Handlers
toggleServerSelection,
batchToggleServers,
getServerStatusIconProps,
// Dialog state
selectedToolForConfig,
isConfigModalOpen,
getConfigDialogProps,
// Utilities
localize,
};
}

View file

@ -863,6 +863,7 @@
"com_ui_mcp_servers": "MCP Servers",
"com_ui_mcp_update_var": "Update {{0}}",
"com_ui_mcp_url": "MCP Server URL",
"com_ui_mcp_init_cancelled": "MCP server '{{0}}' initialization was cancelled due to simultaneous request",
"com_ui_medium": "Medium",
"com_ui_memories": "Memories",
"com_ui_memories_allow_create": "Allow creating Memories",