🔧 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

client/src/components/Chat/Input/BadgeRow.tsx

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

client/src/components/Chat/Input/MCPConfigDialog.tsx

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

client/src/components/Chat/Input/MCPSelect.tsx

@@ -3,8 +3,11 @@ import { MultiSelect, MCPIcon } from '@librechat/client';
 import MCPServerStatusIcon from '~/components/MCP/MCPServerStatusIcon';
 import { useMCPServerManager } from '~/hooks/MCP/useMCPServerManager';
 import MCPConfigDialog from '~/components/MCP/MCPConfigDialog';
+import { useBadgeRowContext } from '~/Providers';
 
-function MCPSelect() {
+type MCPSelectProps = { conversationId?: string | null };
+
+function MCPSelectContent({ conversationId }: MCPSelectProps) {
   const {
     configuredServers,
     mcpValues,
@@ -15,7 +18,7 @@ function MCPSelect() {
     getConfigDialogProps,
     isInitializing,
     localize,
-  } = useMCPServerManager();
+  } = useMCPServerManager({ conversationId });
 
   const renderSelectedValues = useCallback(
     (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"
         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);

client/src/components/Chat/Input/MCPSubMenu.tsx

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

client/src/components/Chat/Input/ToolsDropdown.tsx

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

client/src/components/MCP/MCPConfigDialog.tsx

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

client/src/components/MCP/ServerInitializationSection.tsx

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

client/src/components/SidePanel/MCP/MCPPanel.tsx

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