🔌 feat: MCP Reinitialization and OAuth in UI (#8598)

* ✨ feat: Add connection status endpoint for MCP servers

- Implemented a new endpoint to retrieve the connection status of all MCP servers without disconnecting idle connections.
- Enhanced MCPManager class with a method to get all user-specific connections.

* feat: add silencer arg to loadCustomConfig function to conditionally print config details

- Modified loadCustomConfig to accept a printConfig parameter that allows me to prevent the entire custom config being printed every time it is called

* fix: new status endpoint actually works now, changes to manager.ts to support it

- Updated the connection status endpoint to utilize Maps for app and user connections, rather than incorrectly treating them as objects.
- Introduced a new method + variable in MCPManager to track servers requiring OAuth discovered at startup.
- Stopped OAuth flow from continuing once detected during startup for a new connection

* refactor: Remove hasAuthConfig since we can get that on the frontend without needing to use the endpoint

* feat: Add MCP connection status query and query key for new endpoint

- Introduced a new query hook `useMCPConnectionStatusQuery` to fetch the connection status of MCP servers.
- Added request in data-service
- Defined the API endpoint for retrieving MCP connection status in api-endpoints.ts.
- Defined new types for MCP connection status responses in the types module.
- Added mcpConnectionStatus key

* feat: Enhance MCPSelect component with connection status and server configuration

- Added connection status handling for MCP servers using the new `useMCPConnectionStatusQuery` hook.
- Implemented logic to display appropriate status icons based on connection state and authentication configuration.
- Updated the server selection logic to utilize configured MCP servers from the startup configuration.
- Refactored the rendering of configuration buttons and status indicators for improved user interaction.

* refactor: move MCPConfigDialog to its own  MCP subdir in ui and update import

* refactor: silence loadCustomConfig in status endpoint

* feat: Add optional pluginKey parameter to getUserPluginAuthValue

* feat: Add MCP authentication values endpoint and related queries

- Implemented a new endpoint to check authentication value flags for specific MCP servers, returning boolean indicators for each custom user variable.
- Added a corresponding query hook `useMCPAuthValuesQuery` to fetch authentication values from the frontend.
- Defined the API endpoint for retrieving MCP authentication values in api-endpoints.ts.
- Updated data-service to include a method for fetching MCP authentication values.
- Introduced new types for MCP authentication values responses in the types module.
- Added a new query key for MCP authentication values.

* feat: Localize MCPSelect component status labels and aria attributes

- Updated the MCPSelect component to use localized strings for connection status labels and aria attributes, enhancing accessibility and internationalization support.
- Added new translation keys for various connection states in the translation.json file.

* feat: Implement filtered MCP values selection based on connection status in MCPSelect

- Added a new `filteredSetMCPValues` function to ensure only connected servers are selectable in the MCPSelect component.
- Updated the rendering logic to visually indicate the connection status of servers by adjusting opacity.
- Enhanced accessibility by localizing the aria-label for the configuration button.

* feat: Add CustomUserVarsSection component for managing user variables

- Introduced a new `CustomUserVarsSection` component to allow users to configure custom variables for MCP servers.
- Integrated localization for user interface elements and added new translation keys for variable management.
- Added functionality to save and revoke user variables, with visual indicators for set/unset states.

* feat: Enhance MCPSelect and MCPConfigDialog with improved state management and UI updates

- Integrated `useQueryClient` to refetch queries for tools, authentication values, and connection status upon successful plugin updates in MCPSelect.
- Simplified plugin key handling by directly using the formatted plugin key in save and revoke operations.
- Updated MCPConfigDialog to include server status indicators and improved dialog content structure for better user experience.
- Added new translation key for active status in the localization files.

* feat: Enhance MCPConfigDialog with dynamic server status badges and localization updates

- Added a helper function to render status badges based on the connection state of the MCP server, improving user feedback on connection status.
- Updated the localization files to include new translation keys for connection states such as "Connecting" and "Offline".
- Refactored the dialog to utilize the new status rendering function for better code organization and readability.

* feat: Implement OAuth handling and server initialization in MCP reinitialize flow

- Added OAuth handling to the MCP reinitialize endpoint, allowing the server to capture and return OAuth URLs when required.
- Updated the MCPConfigDialog to include a new ServerInitializationSection for managing server initialization and OAuth flow.
- Enhanced the user experience by providing feedback on server status and OAuth requirements through localized messages.
- Introduced new translation keys for OAuth-related messages in the localization files.
- Refactored the MCPSelect component to remove unused authentication configuration props.

* feat: Make OAuth actually work / update after OAuth link authorized

- Improved the handling of OAuth flows in the MCP reinitialize process, allowing for immediate return when OAuth is initiated.
- Updated the UserController to extract server names from plugin keys for better logging and connection management.
- Enhanced the MCPSelect component to reflect authentication status based on OAuth requirements.
- Implemented polling for OAuth completion in the ServerInitializationSection to improve user feedback during the connection process.
- Refactored MCPManager to support new OAuth flow initiation logic and connection handling.

* refactor: Simplify MCPPanel component and enhance server status display

- Removed unused imports and state management related to user plugins and server reinitialization.
- Integrated connection status handling directly into the MCPPanel for improved user feedback.
- Updated the rendering logic to display server connection states with visual indicators.
- Refactored the editing view to utilize new components for server initialization and custom user variables management.

* chore: remove comments

* chore: remove unused translation key for MCP panel

* refactor: Rename returnOnOAuthInitiated to returnOnOAuth for clarity

* refactor: attempt initialize on server click

* feat: add cancel OAuth flow functionality and related UI updates

* refactor: move server status icon logic into its own component

* chore: remove old localization strings (makes more sense for icon labels to just use configure stirng since thats where it leads to)

* fix: fix accessibility issues with MCPSelect

* fix: add missing save/revoke mutation logic to MCPPanel

* styling: add margin to checkmark in MultiSelect

* fix: add back in customUserVars check to hide gear config icon for servers without customUserVars

---------

Co-authored-by: Dustin Healy <dustinhealy1@gmail.com>
Co-authored-by: Dustin Healy <54083382+dustinhealy@users.noreply.github.com>

client/src/hooks/MCP/index.ts

@@ -0,0 +1 @@
+export { useMCPServerInitialization } from './useMCPServerInitialization';

client/src/hooks/MCP/useMCPServerInitialization.ts

@@ -0,0 +1,289 @@
+import { useCallback, useState, useEffect, useMemo } from 'react';
+import { useQueryClient } from '@tanstack/react-query';
+import { QueryKeys } from 'librechat-data-provider';
+import {
+  useReinitializeMCPServerMutation,
+  useCancelMCPOAuthMutation,
+} from 'librechat-data-provider/react-query';
+import { useMCPConnectionStatusQuery } from '~/data-provider/Tools/queries';
+import { useToastContext } from '~/Providers';
+import { useLocalize } from '~/hooks';
+import { logger } from '~/utils';
+
+interface UseMCPServerInitializationOptions {
+  onSuccess?: (serverName: string) => void;
+  onOAuthStarted?: (serverName: string, oauthUrl: string) => void;
+  onError?: (serverName: string, error: any) => void;
+}
+
+export function useMCPServerInitialization(options?: UseMCPServerInitializationOptions) {
+  const localize = useLocalize();
+  const { showToast } = useToastContext();
+  const queryClient = useQueryClient();
+
+  // OAuth state management
+  const [oauthPollingServers, setOauthPollingServers] = useState<Map<string, string>>(new Map());
+  const [oauthStartTimes, setOauthStartTimes] = useState<Map<string, number>>(new Map());
+  const [initializingServers, setInitializingServers] = useState<Set<string>>(new Set());
+  const [cancellableServers, setCancellableServers] = useState<Set<string>>(new Set());
+
+  // Get connection status
+  const { data: connectionStatusData } = useMCPConnectionStatusQuery();
+  const connectionStatus = useMemo(
+    () => connectionStatusData?.connectionStatus || {},
+    [connectionStatusData],
+  );
+
+  // Main initialization mutation
+  const reinitializeMutation = useReinitializeMCPServerMutation();
+
+  // Cancel OAuth mutation
+  const cancelOAuthMutation = useCancelMCPOAuthMutation();
+
+  // Helper function to clean up OAuth state
+  const cleanupOAuthState = useCallback((serverName: string) => {
+    setOauthPollingServers((prev) => {
+      const newMap = new Map(prev);
+      newMap.delete(serverName);
+      return newMap;
+    });
+
+    setOauthStartTimes((prev) => {
+      const newMap = new Map(prev);
+      newMap.delete(serverName);
+      return newMap;
+    });
+
+    setInitializingServers((prev) => {
+      const newSet = new Set(prev);
+      newSet.delete(serverName);
+      return newSet;
+    });
+
+    setCancellableServers((prev) => {
+      const newSet = new Set(prev);
+      newSet.delete(serverName);
+      return newSet;
+    });
+  }, []);
+
+  // Cancel OAuth flow
+  const cancelOAuthFlow = useCallback(
+    (serverName: string) => {
+      logger.info(`[MCP OAuth] User cancelling OAuth flow for ${serverName}`);
+
+      cancelOAuthMutation.mutate(serverName, {
+        onSuccess: () => {
+          cleanupOAuthState(serverName);
+          showToast({
+            message: localize('com_ui_mcp_oauth_cancelled', { 0: serverName }),
+            status: 'info',
+          });
+        },
+        onError: (error) => {
+          logger.error(`[MCP OAuth] Failed to cancel OAuth flow for ${serverName}:`, error);
+          // Clean up state anyway
+          cleanupOAuthState(serverName);
+        },
+      });
+    },
+    [cancelOAuthMutation, cleanupOAuthState, showToast, localize],
+  );
+
+  // Helper function to handle successful connection
+  const handleSuccessfulConnection = useCallback(
+    async (serverName: string, message: string) => {
+      showToast({ message, status: 'success' });
+
+      // Force immediate refetch to update UI
+      await Promise.all([
+        queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]),
+        queryClient.refetchQueries([QueryKeys.tools]),
+      ]);
+
+      // Clean up OAuth state
+      cleanupOAuthState(serverName);
+
+      // Call optional success callback
+      options?.onSuccess?.(serverName);
+    },
+    [showToast, queryClient, options, cleanupOAuthState],
+  );
+
+  // Helper function to handle OAuth timeout/failure
+  const handleOAuthFailure = useCallback(
+    (serverName: string, isTimeout: boolean) => {
+      logger.warn(
+        `[MCP OAuth] OAuth ${isTimeout ? 'timed out' : 'failed'} for ${serverName}, stopping poll`,
+      );
+
+      // Clean up OAuth state
+      cleanupOAuthState(serverName);
+
+      // Show error toast
+      showToast({
+        message: isTimeout
+          ? localize('com_ui_mcp_oauth_timeout', { 0: serverName })
+          : localize('com_ui_mcp_init_failed'),
+        status: 'error',
+      });
+    },
+    [showToast, localize, cleanupOAuthState],
+  );
+
+  // Poll for OAuth completion
+  useEffect(() => {
+    if (oauthPollingServers.size === 0) {
+      return;
+    }
+
+    const pollInterval = setInterval(() => {
+      // Check each polling server
+      oauthPollingServers.forEach((oauthUrl, serverName) => {
+        const serverStatus = connectionStatus[serverName];
+
+        // Check for client-side timeout (3 minutes)
+        const startTime = oauthStartTimes.get(serverName);
+        const hasTimedOut = startTime && Date.now() - startTime > 180000; // 3 minutes
+
+        if (serverStatus?.connectionState === 'connected') {
+          // OAuth completed successfully
+          handleSuccessfulConnection(
+            serverName,
+            localize('com_ui_mcp_authenticated_success', { 0: serverName }),
+          );
+        } else if (serverStatus?.connectionState === 'error' || hasTimedOut) {
+          // OAuth failed or timed out
+          handleOAuthFailure(serverName, !!hasTimedOut);
+        }
+
+        setCancellableServers((prev) => new Set(prev).add(serverName));
+      });
+
+      queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]);
+    }, 3500);
+
+    return () => {
+      clearInterval(pollInterval);
+    };
+  }, [
+    oauthPollingServers,
+    oauthStartTimes,
+    connectionStatus,
+    queryClient,
+    handleSuccessfulConnection,
+    handleOAuthFailure,
+    localize,
+  ]);
+
+  // Initialize server function
+  const initializeServer = useCallback(
+    (serverName: string) => {
+      // Prevent spam - check if already initializing
+      if (initializingServers.has(serverName)) {
+        return;
+      }
+
+      // Add to initializing set
+      setInitializingServers((prev) => new Set(prev).add(serverName));
+
+      // Trigger initialization
+      reinitializeMutation.mutate(serverName, {
+        onSuccess: (response: any) => {
+          if (response.success) {
+            if (response.oauthRequired && response.oauthUrl) {
+              // OAuth required - store URL and start polling
+              setOauthPollingServers((prev) => new Map(prev).set(serverName, response.oauthUrl));
+
+              // Track when OAuth started for timeout detection
+              setOauthStartTimes((prev) => new Map(prev).set(serverName, Date.now()));
+
+              // Call optional OAuth callback or open URL directly
+              if (options?.onOAuthStarted) {
+                options.onOAuthStarted(serverName, response.oauthUrl);
+              } else {
+                window.open(response.oauthUrl, '_blank', 'noopener,noreferrer');
+              }
+
+              showToast({
+                message: localize('com_ui_connecting'),
+                status: 'info',
+              });
+            } else if (response.oauthRequired) {
+              // OAuth required but no URL - shouldn't happen
+              showToast({
+                message: localize('com_ui_mcp_oauth_no_url'),
+                status: 'warning',
+              });
+              // Remove from initializing since it failed
+              setInitializingServers((prev) => {
+                const newSet = new Set(prev);
+                newSet.delete(serverName);
+                return newSet;
+              });
+            } else {
+              // Successful connection without OAuth
+              handleSuccessfulConnection(
+                serverName,
+                response.message || localize('com_ui_mcp_initialized_success', { 0: serverName }),
+              );
+            }
+          } else {
+            // Remove from initializing if not successful
+            setInitializingServers((prev) => {
+              const newSet = new Set(prev);
+              newSet.delete(serverName);
+              return newSet;
+            });
+          }
+        },
+        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;
+          });
+          // Call optional error callback
+          options?.onError?.(serverName, error);
+        },
+      });
+    },
+    [
+      reinitializeMutation,
+      showToast,
+      localize,
+      handleSuccessfulConnection,
+      initializingServers,
+      options,
+    ],
+  );
+
+  return {
+    initializeServer,
+    isInitializing: (serverName: string) => initializingServers.has(serverName),
+    isCancellable: (serverName: string) => cancellableServers.has(serverName),
+    initializingServers,
+    oauthPollingServers,
+    oauthStartTimes,
+    connectionStatus,
+    isLoading: reinitializeMutation.isLoading,
+    cancelOAuthFlow,
+  };
+}