🔧 fix: MCP Selection Persist and UI Flicker Issues (#9324)

* refactor: useMCPSelect

    - Add useGetMCPTools to use in useMCPSelect and elsewhere hooks for fetching MCP tools
    - remove memoized key
    - remove use of `useChatContext` and require conversationId as prop

* feat: Add MCPPanelContext and integrate conversationId as prop for useMCPSelect across components

- Introduced MCPPanelContext to manage conversationId state.
- Updated MCPSelect, MCPSubMenu, and MCPConfigDialog to accept conversationId as a prop.
- Modified ToolsDropdown and BadgeRow to pass conversationId to relevant components.
- Refactored MCPPanel to utilize MCPPanelProvider for context management.

* fix: remove nested ternary in ServerInitializationSection

- Replaced conditional operator with if-else statements for better readability in determining button text based on server initialization state and reinitialization status.

* refactor: wrap setValueWrap in useCallback for performance optimization

* refactor: streamline useMCPSelect by consolidating storageKey definition

* fix: prevent clearing selections on page refresh by tracking initial load completion

* refactor: simplify concern of useMCPSelect hook

* refactor: move ConfigFieldDetail interface to common types for better reusability, isolate usage of `useGetMCPTools`

* refactor: integrate mcpServerNames into BadgeRowContext and update ToolsDropdown and MCPSelect components
This commit is contained in:
Danny Avila 2025-08-28 00:44:49 -04:00 committed by GitHub
parent 2483623c88
commit c0511b9a5f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 270 additions and 206 deletions

View file

@ -1,12 +1,19 @@
import React, { createContext, useContext, useEffect, useRef } from 'react'; import React, { createContext, useContext, useEffect, useMemo, useRef } from 'react';
import { useSetRecoilState } from 'recoil'; import { useSetRecoilState } from 'recoil';
import { Tools, Constants, LocalStorageKeys, AgentCapabilities } from 'librechat-data-provider'; import { Tools, Constants, LocalStorageKeys, AgentCapabilities } from 'librechat-data-provider';
import type { TAgentsEndpoint } from 'librechat-data-provider'; import type { TAgentsEndpoint } from 'librechat-data-provider';
import { useSearchApiKeyForm, useGetAgentsConfig, useCodeApiKeyForm, useToolToggle } from '~/hooks'; import {
useSearchApiKeyForm,
useGetAgentsConfig,
useCodeApiKeyForm,
useGetMCPTools,
useToolToggle,
} from '~/hooks';
import { ephemeralAgentByConvoId } from '~/store'; import { ephemeralAgentByConvoId } from '~/store';
interface BadgeRowContextType { interface BadgeRowContextType {
conversationId?: string | null; conversationId?: string | null;
mcpServerNames?: string[] | null;
agentsConfig?: TAgentsEndpoint | null; agentsConfig?: TAgentsEndpoint | null;
webSearch: ReturnType<typeof useToolToggle>; webSearch: ReturnType<typeof useToolToggle>;
artifacts: ReturnType<typeof useToolToggle>; artifacts: ReturnType<typeof useToolToggle>;
@ -37,10 +44,12 @@ export default function BadgeRowProvider({
isSubmitting, isSubmitting,
conversationId, conversationId,
}: BadgeRowProviderProps) { }: BadgeRowProviderProps) {
const hasInitializedRef = useRef(false);
const lastKeyRef = useRef<string>(''); const lastKeyRef = useRef<string>('');
const hasInitializedRef = useRef(false);
const { mcpToolDetails } = useGetMCPTools();
const { agentsConfig } = useGetAgentsConfig(); const { agentsConfig } = useGetAgentsConfig();
const key = conversationId ?? Constants.NEW_CONVO; const key = conversationId ?? Constants.NEW_CONVO;
const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(key)); const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(key));
/** Initialize ephemeralAgent from localStorage on mount and when conversation changes */ /** Initialize ephemeralAgent from localStorage on mount and when conversation changes */
@ -156,11 +165,16 @@ export default function BadgeRowProvider({
isAuthenticated: true, isAuthenticated: true,
}); });
const mcpServerNames = useMemo(() => {
return (mcpToolDetails ?? []).map((tool) => tool.name);
}, [mcpToolDetails]);
const value: BadgeRowContextType = { const value: BadgeRowContextType = {
webSearch, webSearch,
artifacts, artifacts,
fileSearch, fileSearch,
agentsConfig, agentsConfig,
mcpServerNames,
conversationId, conversationId,
codeApiKeyForm, codeApiKeyForm,
codeInterpreter, codeInterpreter,

View file

@ -0,0 +1,31 @@
import React, { createContext, useContext, useMemo } from 'react';
import { Constants } from 'librechat-data-provider';
import { useChatContext } from './ChatContext';
interface MCPPanelContextValue {
conversationId: string;
}
const MCPPanelContext = createContext<MCPPanelContextValue | undefined>(undefined);
export function MCPPanelProvider({ children }: { children: React.ReactNode }) {
const { conversation } = useChatContext();
/** Context value only created when conversationId changes */
const contextValue = useMemo<MCPPanelContextValue>(
() => ({
conversationId: conversation?.conversationId ?? Constants.NEW_CONVO,
}),
[conversation?.conversationId],
);
return <MCPPanelContext.Provider value={contextValue}>{children}</MCPPanelContext.Provider>;
}
export function useMCPPanelContext() {
const context = useContext(MCPPanelContext);
if (!context) {
throw new Error('useMCPPanelContext must be used within MCPPanelProvider');
}
return context;
}

View file

@ -23,6 +23,7 @@ export * from './SetConvoContext';
export * from './SearchContext'; export * from './SearchContext';
export * from './BadgeRowContext'; export * from './BadgeRowContext';
export * from './SidePanelContext'; export * from './SidePanelContext';
export * from './MCPPanelContext';
export * from './ArtifactsContext'; export * from './ArtifactsContext';
export * from './PromptGroupsContext'; export * from './PromptGroupsContext';
export { default as BadgeRowProvider } from './BadgeRowContext'; export { default as BadgeRowProvider } from './BadgeRowContext';

View file

@ -8,6 +8,11 @@ import type * as t from 'librechat-data-provider';
import type { LucideIcon } from 'lucide-react'; import type { LucideIcon } from 'lucide-react';
import type { TranslationKeys } from '~/hooks'; import type { TranslationKeys } from '~/hooks';
export interface ConfigFieldDetail {
title: string;
description: string;
}
export type CodeBarProps = { export type CodeBarProps = {
lang: string; lang: string;
error?: boolean; error?: boolean;

View file

@ -368,7 +368,7 @@ function BadgeRow({
<CodeInterpreter /> <CodeInterpreter />
<FileSearch /> <FileSearch />
<Artifacts /> <Artifacts />
<MCPSelect /> <MCPSelect conversationId={conversationId} />
</> </>
)} )}
{ghostBadge && ( {ghostBadge && (

View file

@ -1,13 +1,9 @@
import React, { useEffect } from 'react'; import React, { useEffect } from 'react';
import { useForm, Controller } from 'react-hook-form'; import { useForm, Controller } from 'react-hook-form';
import { Button, Input, Label, OGDialog, OGDialogTemplate } from '@librechat/client'; import { Button, Input, Label, OGDialog, OGDialogTemplate } from '@librechat/client';
import type { ConfigFieldDetail } from '~/common';
import { useLocalize } from '~/hooks'; import { useLocalize } from '~/hooks';
export interface ConfigFieldDetail {
title: string;
description: string;
}
interface MCPConfigDialogProps { interface MCPConfigDialogProps {
isOpen: boolean; isOpen: boolean;
onOpenChange: (isOpen: boolean) => void; onOpenChange: (isOpen: boolean) => void;
@ -34,7 +30,7 @@ export default function MCPConfigDialog({
control, control,
handleSubmit, handleSubmit,
reset, reset,
formState: { errors, _ }, formState: { errors },
} = useForm<Record<string, string>>({ } = useForm<Record<string, string>>({
defaultValues: initialValues, defaultValues: initialValues,
}); });
@ -56,14 +52,12 @@ export default function MCPConfigDialog({
}; };
const dialogTitle = localize('com_ui_configure_mcp_variables_for', { 0: serverName }); const dialogTitle = localize('com_ui_configure_mcp_variables_for', { 0: serverName });
const dialogDescription = localize('com_ui_mcp_dialog_desc');
return ( return (
<OGDialog open={isOpen} onOpenChange={onOpenChange}> <OGDialog open={isOpen} onOpenChange={onOpenChange}>
<OGDialogTemplate <OGDialogTemplate
className="sm:max-w-lg" className="sm:max-w-lg"
title={dialogTitle} title={dialogTitle}
description={dialogDescription}
headerClassName="px-6 pt-6 pb-4" headerClassName="px-6 pt-6 pb-4"
main={ main={
<form onSubmit={handleSubmit(onFormSubmit)} className="space-y-4 px-6 pb-2"> <form onSubmit={handleSubmit(onFormSubmit)} className="space-y-4 px-6 pb-2">

View file

@ -3,8 +3,11 @@ import { MultiSelect, MCPIcon } from '@librechat/client';
import MCPServerStatusIcon from '~/components/MCP/MCPServerStatusIcon'; import MCPServerStatusIcon from '~/components/MCP/MCPServerStatusIcon';
import { useMCPServerManager } from '~/hooks/MCP/useMCPServerManager'; import { useMCPServerManager } from '~/hooks/MCP/useMCPServerManager';
import MCPConfigDialog from '~/components/MCP/MCPConfigDialog'; import MCPConfigDialog from '~/components/MCP/MCPConfigDialog';
import { useBadgeRowContext } from '~/Providers';
function MCPSelect() { type MCPSelectProps = { conversationId?: string | null };
function MCPSelectContent({ conversationId }: MCPSelectProps) {
const { const {
configuredServers, configuredServers,
mcpValues, mcpValues,
@ -15,7 +18,7 @@ function MCPSelect() {
getConfigDialogProps, getConfigDialogProps,
isInitializing, isInitializing,
localize, localize,
} = useMCPServerManager(); } = useMCPServerManager({ conversationId });
const renderSelectedValues = useCallback( const renderSelectedValues = useCallback(
(values: string[], placeholder?: string) => { (values: string[], placeholder?: string) => {
@ -93,9 +96,17 @@ function MCPSelect() {
selectItemsClassName="border border-blue-600/50 bg-blue-500/10 hover:bg-blue-700/10" 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" 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"
/> />
{configDialogProps && <MCPConfigDialog {...configDialogProps} />} {configDialogProps && (
<MCPConfigDialog {...configDialogProps} conversationId={conversationId} />
)}
</> </>
); );
} }
function MCPSelect(props: MCPSelectProps) {
const { mcpServerNames } = useBadgeRowContext();
if ((mcpServerNames?.length ?? 0) === 0) return null;
return <MCPSelectContent {...props} />;
}
export default memo(MCPSelect); export default memo(MCPSelect);

View file

@ -9,10 +9,11 @@ import { cn } from '~/utils';
interface MCPSubMenuProps { interface MCPSubMenuProps {
placeholder?: string; placeholder?: string;
conversationId?: string | null;
} }
const MCPSubMenu = React.forwardRef<HTMLDivElement, MCPSubMenuProps>( const MCPSubMenu = React.forwardRef<HTMLDivElement, MCPSubMenuProps>(
({ placeholder, ...props }, ref) => { ({ placeholder, conversationId, ...props }, ref) => {
const { const {
configuredServers, configuredServers,
mcpValues, mcpValues,
@ -23,7 +24,7 @@ const MCPSubMenu = React.forwardRef<HTMLDivElement, MCPSubMenuProps>(
getServerStatusIconProps, getServerStatusIconProps,
getConfigDialogProps, getConfigDialogProps,
isInitializing, isInitializing,
} = useMCPServerManager(); } = useMCPServerManager({ conversationId });
const menuStore = Ariakit.useMenuStore({ const menuStore = Ariakit.useMenuStore({
focusLoop: true, focusLoop: true,

View file

@ -10,12 +10,12 @@ import {
PermissionTypes, PermissionTypes,
defaultAgentCapabilities, defaultAgentCapabilities,
} from 'librechat-data-provider'; } from 'librechat-data-provider';
import { useLocalize, useHasAccess, useAgentCapabilities, useMCPSelect } from '~/hooks'; import { useLocalize, useHasAccess, useAgentCapabilities } from '~/hooks';
import ArtifactsSubMenu from '~/components/Chat/Input/ArtifactsSubMenu'; import ArtifactsSubMenu from '~/components/Chat/Input/ArtifactsSubMenu';
import MCPSubMenu from '~/components/Chat/Input/MCPSubMenu'; import MCPSubMenu from '~/components/Chat/Input/MCPSubMenu';
import { useGetStartupConfig } from '~/data-provider';
import { useBadgeRowContext } from '~/Providers'; import { useBadgeRowContext } from '~/Providers';
import { cn } from '~/utils'; import { cn } from '~/utils';
import { useGetStartupConfig } from '~/data-provider';
interface ToolsDropdownProps { interface ToolsDropdownProps {
disabled?: boolean; disabled?: boolean;
@ -30,11 +30,12 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => {
artifacts, artifacts,
fileSearch, fileSearch,
agentsConfig, agentsConfig,
mcpServerNames,
conversationId,
codeApiKeyForm, codeApiKeyForm,
codeInterpreter, codeInterpreter,
searchApiKeyForm, searchApiKeyForm,
} = useBadgeRowContext(); } = useBadgeRowContext();
const mcpSelect = useMCPSelect();
const { data: startupConfig } = useGetStartupConfig(); const { data: startupConfig } = useGetStartupConfig();
const { codeEnabled, webSearchEnabled, artifactsEnabled, fileSearchEnabled } = const { codeEnabled, webSearchEnabled, artifactsEnabled, fileSearchEnabled } =
@ -56,7 +57,6 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => {
} = codeInterpreter; } = codeInterpreter;
const { isPinned: isFileSearchPinned, setIsPinned: setIsFileSearchPinned } = fileSearch; const { isPinned: isFileSearchPinned, setIsPinned: setIsFileSearchPinned } = fileSearch;
const { isPinned: isArtifactsPinned, setIsPinned: setIsArtifactsPinned } = artifacts; const { isPinned: isArtifactsPinned, setIsPinned: setIsArtifactsPinned } = artifacts;
const { mcpServerNames } = mcpSelect;
const canUseWebSearch = useHasAccess({ const canUseWebSearch = useHasAccess({
permissionType: PermissionTypes.WEB_SEARCH, permissionType: PermissionTypes.WEB_SEARCH,
@ -290,7 +290,9 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => {
if (mcpServerNames && mcpServerNames.length > 0) { if (mcpServerNames && mcpServerNames.length > 0) {
dropdownItems.push({ dropdownItems.push({
hideOnClick: false, hideOnClick: false,
render: (props) => <MCPSubMenu {...props} placeholder={mcpPlaceholder} />, render: (props) => (
<MCPSubMenu {...props} placeholder={mcpPlaceholder} conversationId={conversationId} />
),
}); });
} }

View file

@ -8,15 +8,11 @@ import {
OGDialogContent, OGDialogContent,
} from '@librechat/client'; } from '@librechat/client';
import type { MCPServerStatus } from 'librechat-data-provider'; import type { MCPServerStatus } from 'librechat-data-provider';
import type { ConfigFieldDetail } from '~/common';
import ServerInitializationSection from './ServerInitializationSection'; import ServerInitializationSection from './ServerInitializationSection';
import CustomUserVarsSection from './CustomUserVarsSection'; import CustomUserVarsSection from './CustomUserVarsSection';
import { useLocalize } from '~/hooks'; import { useLocalize } from '~/hooks';
export interface ConfigFieldDetail {
title: string;
description: string;
}
interface MCPConfigDialogProps { interface MCPConfigDialogProps {
isOpen: boolean; isOpen: boolean;
onOpenChange: (isOpen: boolean) => void; onOpenChange: (isOpen: boolean) => void;
@ -27,6 +23,7 @@ interface MCPConfigDialogProps {
onRevoke?: () => void; onRevoke?: () => void;
serverName: string; serverName: string;
serverStatus?: MCPServerStatus; serverStatus?: MCPServerStatus;
conversationId?: string | null;
} }
export default function MCPConfigDialog({ export default function MCPConfigDialog({
@ -38,6 +35,7 @@ export default function MCPConfigDialog({
onRevoke, onRevoke,
serverName, serverName,
serverStatus, serverStatus,
conversationId,
}: MCPConfigDialogProps) { }: MCPConfigDialogProps) {
const localize = useLocalize(); const localize = useLocalize();
@ -126,6 +124,7 @@ export default function MCPConfigDialog({
{/* Server Initialization Section */} {/* Server Initialization Section */}
<ServerInitializationSection <ServerInitializationSection
serverName={serverName} serverName={serverName}
conversationId={conversationId}
requiresOAuth={serverStatus?.requiresOAuth || false} requiresOAuth={serverStatus?.requiresOAuth || false}
hasCustomUserVars={fieldsSchema && Object.keys(fieldsSchema).length > 0} hasCustomUserVars={fieldsSchema && Object.keys(fieldsSchema).length > 0}
/> />

View file

@ -9,12 +9,14 @@ interface ServerInitializationSectionProps {
serverName: string; serverName: string;
requiresOAuth: boolean; requiresOAuth: boolean;
hasCustomUserVars?: boolean; hasCustomUserVars?: boolean;
conversationId?: string | null;
} }
export default function ServerInitializationSection({ export default function ServerInitializationSection({
sidePanel = false,
serverName, serverName,
requiresOAuth, requiresOAuth,
conversationId,
sidePanel = false,
hasCustomUserVars = false, hasCustomUserVars = false,
}: ServerInitializationSectionProps) { }: ServerInitializationSectionProps) {
const localize = useLocalize(); const localize = useLocalize();
@ -26,7 +28,7 @@ export default function ServerInitializationSection({
isInitializing, isInitializing,
isCancellable, isCancellable,
getOAuthUrl, getOAuthUrl,
} = useMCPServerManager(); } = useMCPServerManager({ conversationId });
const serverStatus = connectionStatus[serverName]; const serverStatus = connectionStatus[serverName];
const isConnected = serverStatus?.connectionState === 'connected'; const isConnected = serverStatus?.connectionState === 'connected';
@ -69,13 +71,18 @@ export default function ServerInitializationSection({
const isReinit = shouldShowReinit; const isReinit = shouldShowReinit;
const outerClass = isReinit ? 'flex justify-start' : 'flex justify-end'; const outerClass = isReinit ? 'flex justify-start' : 'flex justify-end';
const buttonVariant = isReinit ? undefined : 'default'; const buttonVariant = isReinit ? undefined : 'default';
const buttonText = isServerInitializing
? localize('com_ui_loading') let buttonText = '';
: isReinit if (isServerInitializing) {
? localize('com_ui_reinitialize') buttonText = localize('com_ui_loading');
: requiresOAuth } else if (isReinit) {
? localize('com_ui_authenticate') buttonText = localize('com_ui_reinitialize');
: localize('com_ui_mcp_initialize'); } else if (requiresOAuth) {
buttonText = localize('com_ui_authenticate');
} else {
buttonText = localize('com_ui_mcp_initialize');
}
const icon = isServerInitializing ? ( const icon = isServerInitializing ? (
<Spinner className="h-4 w-4" /> <Spinner className="h-4 w-4" />
) : ( ) : (

View file

@ -8,15 +8,16 @@ import type { TUpdateUserPlugins } from 'librechat-data-provider';
import ServerInitializationSection from '~/components/MCP/ServerInitializationSection'; import ServerInitializationSection from '~/components/MCP/ServerInitializationSection';
import { useMCPConnectionStatusQuery } from '~/data-provider/Tools/queries'; import { useMCPConnectionStatusQuery } from '~/data-provider/Tools/queries';
import CustomUserVarsSection from '~/components/MCP/CustomUserVarsSection'; import CustomUserVarsSection from '~/components/MCP/CustomUserVarsSection';
import BadgeRowProvider from '~/Providers/BadgeRowContext'; import { MCPPanelProvider, useMCPPanelContext } from '~/Providers';
import { useGetStartupConfig } from '~/data-provider'; import { useGetStartupConfig } from '~/data-provider';
import MCPPanelSkeleton from './MCPPanelSkeleton'; import MCPPanelSkeleton from './MCPPanelSkeleton';
import { useLocalize } from '~/hooks'; import { useLocalize } from '~/hooks';
function MCPPanelContent() { function MCPPanelContent() {
const localize = useLocalize(); const localize = useLocalize();
const { showToast } = useToastContext();
const queryClient = useQueryClient(); const queryClient = useQueryClient();
const { showToast } = useToastContext();
const { conversationId } = useMCPPanelContext();
const { data: startupConfig, isLoading: startupConfigLoading } = useGetStartupConfig(); const { data: startupConfig, isLoading: startupConfigLoading } = useGetStartupConfig();
const { data: connectionStatusData } = useMCPConnectionStatusQuery(); const { data: connectionStatusData } = useMCPConnectionStatusQuery();
const [selectedServerNameForEditing, setSelectedServerNameForEditing] = useState<string | null>( const [selectedServerNameForEditing, setSelectedServerNameForEditing] = useState<string | null>(
@ -153,6 +154,7 @@ function MCPPanelContent() {
<ServerInitializationSection <ServerInitializationSection
sidePanel={true} sidePanel={true}
conversationId={conversationId}
serverName={selectedServerNameForEditing} serverName={selectedServerNameForEditing}
requiresOAuth={serverStatus?.requiresOAuth || false} requiresOAuth={serverStatus?.requiresOAuth || false}
hasCustomUserVars={ hasCustomUserVars={
@ -204,8 +206,8 @@ function MCPPanelContent() {
export default function MCPPanel() { export default function MCPPanel() {
return ( return (
<BadgeRowProvider> <MCPPanelProvider>
<MCPPanelContent /> <MCPPanelContent />
</BadgeRowProvider> </MCPPanelProvider>
); );
} }

View file

@ -1 +1,3 @@
export * from './useMCPSelect';
export * from './useGetMCPTools';
export { useMCPServerManager } from './useMCPServerManager'; export { useMCPServerManager } from './useMCPServerManager';

View file

@ -0,0 +1,43 @@
import { useMemo } from 'react';
import { Constants, EModelEndpoint } from 'librechat-data-provider';
import type { TPlugin } from 'librechat-data-provider';
import { useAvailableToolsQuery, useGetStartupConfig } from '~/data-provider';
export function useGetMCPTools() {
const { data: startupConfig } = useGetStartupConfig();
const { data: rawMcpTools } = useAvailableToolsQuery(EModelEndpoint.agents, {
select: (data: TPlugin[]) => {
const mcpToolsMap = new Map<string, TPlugin>();
data.forEach((tool) => {
const isMCP = tool.pluginKey.includes(Constants.mcp_delimiter);
if (isMCP) {
const parts = tool.pluginKey.split(Constants.mcp_delimiter);
const serverName = parts[parts.length - 1];
if (!mcpToolsMap.has(serverName)) {
mcpToolsMap.set(serverName, {
name: serverName,
pluginKey: tool.pluginKey,
authConfig: tool.authConfig,
authenticated: tool.authenticated,
});
}
}
});
return Array.from(mcpToolsMap.values());
},
});
const mcpToolDetails = useMemo(() => {
if (!rawMcpTools || !startupConfig?.mcpServers) {
return rawMcpTools;
}
return rawMcpTools.filter((tool) => {
const serverConfig = startupConfig?.mcpServers?.[tool.name];
return serverConfig?.chatMenu !== false;
});
}, [rawMcpTools, startupConfig?.mcpServers]);
return {
mcpToolDetails,
};
}

View file

@ -0,0 +1,72 @@
import { useRef, useCallback, useMemo } from 'react';
import { useRecoilState } from 'recoil';
import { Constants, LocalStorageKeys } from 'librechat-data-provider';
import useLocalStorage from '~/hooks/useLocalStorageAlt';
import { ephemeralAgentByConvoId } from '~/store';
const storageCondition = (value: unknown, rawCurrentValue?: string | null) => {
if (rawCurrentValue) {
try {
const currentValue = rawCurrentValue?.trim() ?? '';
if (currentValue.length > 2) {
return true;
}
} catch (e) {
console.error(e);
}
}
return Array.isArray(value) && value.length > 0;
};
export function useMCPSelect({ conversationId }: { conversationId?: string | null }) {
const key = conversationId ?? Constants.NEW_CONVO;
const [ephemeralAgent, setEphemeralAgent] = useRecoilState(ephemeralAgentByConvoId(key));
const storageKey = `${LocalStorageKeys.LAST_MCP_}${key}`;
const mcpState = useMemo(() => {
return ephemeralAgent?.mcp ?? [];
}, [ephemeralAgent?.mcp]);
const setSelectedValues = useCallback(
(values: string[] | null | undefined) => {
if (!values) {
return;
}
if (!Array.isArray(values)) {
return;
}
setEphemeralAgent((prev) => ({
...prev,
mcp: values,
}));
},
[setEphemeralAgent],
);
const [mcpValues, setMCPValuesRaw] = useLocalStorage<string[]>(
storageKey,
mcpState,
setSelectedValues,
storageCondition,
);
const setMCPValuesRawRef = useRef(setMCPValuesRaw);
setMCPValuesRawRef.current = setMCPValuesRaw;
/** Create a stable memoized setter to avoid re-creating it on every render and causing an infinite render loop */
const setMCPValues = useCallback((value: string[]) => {
setMCPValuesRawRef.current(value);
}, []);
const [isPinned, setIsPinned] = useLocalStorage<boolean>(
`${LocalStorageKeys.PIN_MCP_}${key}`,
true,
);
return {
isPinned,
mcpValues,
setIsPinned,
setMCPValues,
};
}

View file

@ -8,10 +8,10 @@ import {
useReinitializeMCPServerMutation, useReinitializeMCPServerMutation,
} from 'librechat-data-provider/react-query'; } from 'librechat-data-provider/react-query';
import type { TUpdateUserPlugins, TPlugin } from 'librechat-data-provider'; import type { TUpdateUserPlugins, TPlugin } from 'librechat-data-provider';
import type { ConfigFieldDetail } from '~/components/MCP/MCPConfigDialog'; import type { ConfigFieldDetail } from '~/common';
import { useMCPConnectionStatusQuery } from '~/data-provider/Tools/queries'; import { useMCPConnectionStatusQuery } from '~/data-provider/Tools/queries';
import { useLocalize, useMCPSelect, useGetMCPTools } from '~/hooks';
import { useGetStartupConfig } from '~/data-provider'; import { useGetStartupConfig } from '~/data-provider';
import { useLocalize, useMCPSelect } from '~/hooks';
interface ServerState { interface ServerState {
isInitializing: boolean; isInitializing: boolean;
@ -21,13 +21,14 @@ interface ServerState {
pollInterval: NodeJS.Timeout | null; pollInterval: NodeJS.Timeout | null;
} }
export function useMCPServerManager() { export function useMCPServerManager({ conversationId }: { conversationId?: string | null }) {
const localize = useLocalize(); const localize = useLocalize();
const { showToast } = useToastContext();
const mcpSelect = useMCPSelect();
const { data: startupConfig } = useGetStartupConfig();
const { mcpValues, setMCPValues, mcpToolDetails, isPinned, setIsPinned } = mcpSelect;
const queryClient = useQueryClient(); const queryClient = useQueryClient();
const { showToast } = useToastContext();
const { mcpToolDetails } = useGetMCPTools();
const mcpSelect = useMCPSelect({ conversationId });
const { data: startupConfig } = useGetStartupConfig();
const { mcpValues, setMCPValues, isPinned, setIsPinned } = mcpSelect;
const [isConfigModalOpen, setIsConfigModalOpen] = useState(false); const [isConfigModalOpen, setIsConfigModalOpen] = useState(false);
const [selectedToolForConfig, setSelectedToolForConfig] = useState<TPlugin | null>(null); const [selectedToolForConfig, setSelectedToolForConfig] = useState<TPlugin | null>(null);
@ -90,7 +91,21 @@ export function useMCPServerManager() {
[connectionStatusData?.connectionStatus], [connectionStatusData?.connectionStatus],
); );
/** Filter disconnected servers when values change, but only after initial load
This prevents clearing selections on page refresh when servers haven't connected yet
*/
const hasInitialLoadCompleted = useRef(false);
useEffect(() => { useEffect(() => {
if (!connectionStatusData || Object.keys(connectionStatus).length === 0) {
return;
}
if (!hasInitialLoadCompleted.current) {
hasInitialLoadCompleted.current = true;
return;
}
if (!mcpValues?.length) return; if (!mcpValues?.length) return;
const connectedSelected = mcpValues.filter( const connectedSelected = mcpValues.filter(
@ -100,7 +115,7 @@ export function useMCPServerManager() {
if (connectedSelected.length !== mcpValues.length) { if (connectedSelected.length !== mcpValues.length) {
setMCPValues(connectedSelected); setMCPValues(connectedSelected);
} }
}, [connectionStatus, mcpValues, setMCPValues]); }, [connectionStatus, connectionStatusData, mcpValues, setMCPValues]);
const updateServerState = useCallback((serverName: string, updates: Partial<ServerState>) => { const updateServerState = useCallback((serverName: string, updates: Partial<ServerState>) => {
setServerStates((prev) => { setServerStates((prev) => {
@ -486,12 +501,12 @@ export function useMCPServerManager() {
}; };
}, },
[ [
isCancellable,
mcpToolDetails, mcpToolDetails,
isInitializing,
cancelOAuthFlow,
connectionStatus, connectionStatus,
startupConfig?.mcpServers, startupConfig?.mcpServers,
isInitializing,
isCancellable,
cancelOAuthFlow,
], ],
); );
@ -547,7 +562,6 @@ export function useMCPServerManager() {
mcpValues, mcpValues,
setMCPValues, setMCPValues,
mcpToolDetails,
isPinned, isPinned,
setIsPinned, setIsPinned,
placeholderText, placeholderText,

View file

@ -1,4 +1,3 @@
export * from './useMCPSelect';
export * from './useToolToggle'; export * from './useToolToggle';
export { default as useAuthCodeTool } from './useAuthCodeTool'; export { default as useAuthCodeTool } from './useAuthCodeTool';
export { default as usePluginInstall } from './usePluginInstall'; export { default as usePluginInstall } from './usePluginInstall';

View file

@ -1,136 +0,0 @@
import { useRef, useEffect, useCallback, useMemo } from 'react';
import { useRecoilState } from 'recoil';
import { Constants, LocalStorageKeys, EModelEndpoint } from 'librechat-data-provider';
import type { TPlugin } from 'librechat-data-provider';
import { useAvailableToolsQuery, useGetStartupConfig } from '~/data-provider';
import useLocalStorage from '~/hooks/useLocalStorageAlt';
import { ephemeralAgentByConvoId } from '~/store';
import { useChatContext } from '~/Providers';
const storageCondition = (value: unknown, rawCurrentValue?: string | null) => {
if (rawCurrentValue) {
try {
const currentValue = rawCurrentValue?.trim() ?? '';
if (currentValue.length > 2) {
return true;
}
} catch (e) {
console.error(e);
}
}
return Array.isArray(value) && value.length > 0;
};
export function useMCPSelect() {
const { conversation } = useChatContext();
const key = useMemo(
() => conversation?.conversationId ?? Constants.NEW_CONVO,
[conversation?.conversationId],
);
const hasSetFetched = useRef<string | null>(null);
const [ephemeralAgent, setEphemeralAgent] = useRecoilState(ephemeralAgentByConvoId(key));
const { data: startupConfig } = useGetStartupConfig();
const { data: rawMcpTools, isFetched } = useAvailableToolsQuery(EModelEndpoint.agents, {
select: (data: TPlugin[]) => {
const mcpToolsMap = new Map<string, TPlugin>();
data.forEach((tool) => {
const isMCP = tool.pluginKey.includes(Constants.mcp_delimiter);
if (isMCP) {
const parts = tool.pluginKey.split(Constants.mcp_delimiter);
const serverName = parts[parts.length - 1];
if (!mcpToolsMap.has(serverName)) {
mcpToolsMap.set(serverName, {
name: serverName,
pluginKey: tool.pluginKey,
authConfig: tool.authConfig,
authenticated: tool.authenticated,
});
}
}
});
return Array.from(mcpToolsMap.values());
},
});
const mcpToolDetails = useMemo(() => {
if (!rawMcpTools || !startupConfig?.mcpServers) {
return rawMcpTools;
}
return rawMcpTools.filter((tool) => {
const serverConfig = startupConfig?.mcpServers?.[tool.name];
return serverConfig?.chatMenu !== false;
});
}, [rawMcpTools, startupConfig?.mcpServers]);
const mcpState = useMemo(() => {
return ephemeralAgent?.mcp ?? [];
}, [ephemeralAgent?.mcp]);
const setSelectedValues = useCallback(
(values: string[] | null | undefined) => {
if (!values) {
return;
}
if (!Array.isArray(values)) {
return;
}
setEphemeralAgent((prev) => ({
...prev,
mcp: values,
}));
},
[setEphemeralAgent],
);
const [mcpValues, setMCPValuesRaw] = useLocalStorage<string[]>(
`${LocalStorageKeys.LAST_MCP_}${key}`,
mcpState,
setSelectedValues,
storageCondition,
);
const setMCPValuesRawRef = useRef(setMCPValuesRaw);
setMCPValuesRawRef.current = setMCPValuesRaw;
// Create a stable memoized setter to avoid re-creating it on every render and causing an infinite render loop
const setMCPValues = useCallback((value: string[]) => {
setMCPValuesRawRef.current(value);
}, []);
const [isPinned, setIsPinned] = useLocalStorage<boolean>(
`${LocalStorageKeys.PIN_MCP_}${key}`,
true,
);
useEffect(() => {
if (hasSetFetched.current === key) {
return;
}
if (!isFetched) {
return;
}
hasSetFetched.current = key;
if ((mcpToolDetails?.length ?? 0) > 0) {
setMCPValues(mcpValues.filter((mcp) => mcpToolDetails?.some((tool) => tool.name === mcp)));
return;
}
setMCPValues([]);
}, [isFetched, setMCPValues, mcpToolDetails, key, mcpValues]);
const mcpServerNames = useMemo(() => {
return (mcpToolDetails ?? []).map((tool) => tool.name);
}, [mcpToolDetails]);
return {
isPinned,
mcpValues,
setIsPinned,
setMCPValues,
mcpServerNames,
ephemeralAgent,
mcpToolDetails,
setEphemeralAgent,
};
}

View file

@ -5,7 +5,7 @@
* - Also value will be updated everywhere, when value updated (via `storage` event) * - Also value will be updated everywhere, when value updated (via `storage` event)
*/ */
import { useEffect, useState } from 'react'; import { useEffect, useState, useCallback } from 'react';
export default function useLocalStorage<T>( export default function useLocalStorage<T>(
key: string, key: string,
@ -47,7 +47,8 @@ export default function useLocalStorage<T>(
// eslint-disable-next-line react-hooks/exhaustive-deps // eslint-disable-next-line react-hooks/exhaustive-deps
}, [key, globalSetState]); }, [key, globalSetState]);
const setValueWrap = (value: T) => { const setValueWrap = useCallback(
(value: T) => {
try { try {
setValue(value); setValue(value);
const storeLocal = () => { const storeLocal = () => {
@ -63,7 +64,9 @@ export default function useLocalStorage<T>(
} catch (e) { } catch (e) {
console.error(e); console.error(e);
} }
}; },
[key, globalSetState, storageCondition],
);
return [value, setValueWrap]; return [value, setValueWrap];
} }