From 74d8a3824c4fa1caf62265784dc5177ddb52fa2b Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 22 Jul 2025 22:52:45 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=8C=20feat:=20MCP=20Reinitialization?= =?UTF-8?q?=20and=20OAuth=20in=20UI=20(#8598)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ✨ 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 Co-authored-by: Dustin Healy <54083382+dustinhealy@users.noreply.github.com> --- api/server/controllers/UserController.js | 8 +- api/server/routes/mcp.js | 371 ++++++++++++++++-- .../services/Config/loadCustomConfig.js | 10 +- api/server/services/PluginService.js | 21 +- .../src/components/Chat/Input/MCPSelect.tsx | 281 ++++++++++--- .../src/components/SidePanel/MCP/MCPPanel.tsx | 319 ++++++--------- .../ui/MCP/CustomUserVarsSection.tsx | 161 ++++++++ .../src/components/ui/MCP/MCPConfigDialog.tsx | 139 +++++++ .../components/ui/MCP/MCPServerStatusIcon.tsx | 190 +++++++++ .../ui/MCP/ServerInitializationSection.tsx | 131 +++++++ client/src/components/ui/MCPConfigDialog.tsx | 122 ------ client/src/components/ui/MultiSelect.tsx | 2 +- client/src/data-provider/Tools/queries.ts | 33 ++ client/src/hooks/MCP/index.ts | 1 + .../hooks/MCP/useMCPServerInitialization.ts | 289 ++++++++++++++ client/src/locales/en/translation.json | 23 +- packages/api/src/mcp/manager.ts | 91 +++-- packages/data-provider/src/api-endpoints.ts | 8 + packages/data-provider/src/data-service.ts | 12 + packages/data-provider/src/keys.ts | 2 + .../src/react-query/react-query-service.ts | 22 +- packages/data-provider/src/types/mutations.ts | 10 + packages/data-provider/src/types/queries.ts | 16 + 23 files changed, 1812 insertions(+), 450 deletions(-) create mode 100644 client/src/components/ui/MCP/CustomUserVarsSection.tsx create mode 100644 client/src/components/ui/MCP/MCPConfigDialog.tsx create mode 100644 client/src/components/ui/MCP/MCPServerStatusIcon.tsx create mode 100644 client/src/components/ui/MCP/ServerInitializationSection.tsx delete mode 100644 client/src/components/ui/MCPConfigDialog.tsx create mode 100644 client/src/hooks/MCP/index.ts create mode 100644 client/src/hooks/MCP/useMCPServerInitialization.ts diff --git a/api/server/controllers/UserController.js b/api/server/controllers/UserController.js index e570d5ca14..c91b672724 100644 --- a/api/server/controllers/UserController.js +++ b/api/server/controllers/UserController.js @@ -175,14 +175,16 @@ const updateUserPluginsController = async (req, res) => { try { const mcpManager = getMCPManager(user.id); if (mcpManager) { + // Extract server name from pluginKey (format: "mcp_") + const serverName = pluginKey.replace(Constants.mcp_prefix, ''); logger.info( - `[updateUserPluginsController] Disconnecting MCP connections for user ${user.id} after plugin auth update for ${pluginKey}.`, + `[updateUserPluginsController] Disconnecting MCP server ${serverName} for user ${user.id} after plugin auth update for ${pluginKey}.`, ); - await mcpManager.disconnectUserConnections(user.id); + await mcpManager.disconnectUserConnection(user.id, serverName); } } catch (disconnectError) { logger.error( - `[updateUserPluginsController] Error disconnecting MCP connections for user ${user.id} after plugin auth update:`, + `[updateUserPluginsController] Error disconnecting MCP connection for user ${user.id} after plugin auth update:`, disconnectError, ); // Do not fail the request for this, but log it. diff --git a/api/server/routes/mcp.js b/api/server/routes/mcp.js index b9084f982d..5e1db18e67 100644 --- a/api/server/routes/mcp.js +++ b/api/server/routes/mcp.js @@ -120,9 +120,73 @@ router.get('/:serverName/oauth/callback', async (req, res) => { const tokens = await MCPOAuthHandler.completeOAuthFlow(flowId, code, flowManager); logger.info('[MCP OAuth] OAuth flow completed, tokens received in callback route'); - // For system-level OAuth, we need to store the tokens and retry the connection - if (flowState.userId === 'system') { - logger.debug(`[MCP OAuth] System-level OAuth completed for ${serverName}`); + // Try to establish the MCP connection with the new tokens + try { + const mcpManager = getMCPManager(flowState.userId); + logger.debug(`[MCP OAuth] Attempting to reconnect ${serverName} with new OAuth tokens`); + + // For user-level OAuth, try to establish the connection + if (flowState.userId !== 'system') { + // We need to get the user object - in this case we'll need to reconstruct it + const user = { id: flowState.userId }; + + // Try to establish connection with the new tokens + const userConnection = await mcpManager.getUserConnection({ + user, + serverName, + flowManager, + tokenMethods: { + findToken, + updateToken, + createToken, + deleteTokens, + }, + }); + + logger.info( + `[MCP OAuth] Successfully reconnected ${serverName} for user ${flowState.userId}`, + ); + + // Fetch and cache tools now that we have a successful connection + const userTools = (await getCachedTools({ userId: flowState.userId })) || {}; + + // Remove any old tools from this server in the user's cache + const mcpDelimiter = Constants.mcp_delimiter; + for (const key of Object.keys(userTools)) { + if (key.endsWith(`${mcpDelimiter}${serverName}`)) { + delete userTools[key]; + } + } + + // Add the new tools from this server + const tools = await userConnection.fetchTools(); + for (const tool of tools) { + const name = `${tool.name}${Constants.mcp_delimiter}${serverName}`; + userTools[name] = { + type: 'function', + ['function']: { + name, + description: tool.description, + parameters: tool.inputSchema, + }, + }; + } + + // Save the updated user tool cache + await setCachedTools(userTools, { userId: flowState.userId }); + + logger.debug( + `[MCP OAuth] Cached ${tools.length} tools for ${serverName} user ${flowState.userId}`, + ); + } else { + logger.debug(`[MCP OAuth] System-level OAuth completed for ${serverName}`); + } + } catch (error) { + // Don't fail the OAuth callback if reconnection fails - the tokens are still saved + logger.warn( + `[MCP OAuth] Failed to reconnect ${serverName} after OAuth, but tokens are saved:`, + error, + ); } /** ID of the flow that the tool/connection is waiting for */ @@ -205,6 +269,53 @@ router.get('/oauth/status/:flowId', async (req, res) => { } }); +/** + * Cancel OAuth flow + * This endpoint cancels a pending OAuth flow + */ +router.post('/oauth/cancel/:serverName', requireJwtAuth, async (req, res) => { + try { + const { serverName } = req.params; + const user = req.user; + + if (!user?.id) { + return res.status(401).json({ error: 'User not authenticated' }); + } + + logger.info(`[MCP OAuth Cancel] Cancelling OAuth flow for ${serverName} by user ${user.id}`); + + const flowsCache = getLogStores(CacheKeys.FLOWS); + const flowManager = getFlowStateManager(flowsCache); + + // Generate the flow ID for this user/server combination + const flowId = MCPOAuthHandler.generateFlowId(user.id, serverName); + + // Check if flow exists + const flowState = await flowManager.getFlowState(flowId, 'mcp_oauth'); + + if (!flowState) { + logger.debug(`[MCP OAuth Cancel] No active flow found for ${serverName}`); + return res.json({ + success: true, + message: 'No active OAuth flow to cancel', + }); + } + + // Cancel the flow by marking it as failed + await flowManager.completeFlow(flowId, 'mcp_oauth', null, 'User cancelled OAuth flow'); + + logger.info(`[MCP OAuth Cancel] Successfully cancelled OAuth flow for ${serverName}`); + + res.json({ + success: true, + message: `OAuth flow for ${serverName} cancelled successfully`, + }); + } catch (error) { + logger.error('[MCP OAuth Cancel] Failed to cancel OAuth flow', error); + res.status(500).json({ error: 'Failed to cancel OAuth flow' }); + } +}); + /** * Reinitialize MCP server * This endpoint allows reinitializing a specific MCP server @@ -251,6 +362,9 @@ router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => { } let userConnection = null; + let oauthRequired = false; + let oauthUrl = null; + try { userConnection = await mcpManager.getUserConnection({ user, @@ -263,43 +377,87 @@ router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => { createToken, deleteTokens, }, + returnOnOAuth: true, // Return immediately when OAuth is initiated + // Add OAuth handlers to capture the OAuth URL when needed + oauthStart: async (authURL) => { + logger.info(`[MCP Reinitialize] OAuth URL received: ${authURL}`); + oauthUrl = authURL; + oauthRequired = true; + }, }); + + logger.info(`[MCP Reinitialize] Successfully established connection for ${serverName}`); } catch (err) { - logger.error(`[MCP Reinitialize] Error initializing MCP server ${serverName} for user:`, err); - return res.status(500).json({ error: 'Failed to reinitialize MCP server for user' }); - } + logger.info(`[MCP Reinitialize] getUserConnection threw error: ${err.message}`); + logger.info( + `[MCP Reinitialize] OAuth state - oauthRequired: ${oauthRequired}, oauthUrl: ${oauthUrl ? 'present' : 'null'}`, + ); - const userTools = (await getCachedTools({ userId: user.id })) || {}; + // Check if this is an OAuth error - if so, the flow state should be set up now + const isOAuthError = + err.message?.includes('OAuth') || + err.message?.includes('authentication') || + err.message?.includes('401'); - // Remove any old tools from this server in the user's cache - const mcpDelimiter = Constants.mcp_delimiter; - for (const key of Object.keys(userTools)) { - if (key.endsWith(`${mcpDelimiter}${serverName}`)) { - delete userTools[key]; + const isOAuthFlowInitiated = err.message === 'OAuth flow initiated - return early'; + + if (isOAuthError || oauthRequired || isOAuthFlowInitiated) { + logger.info( + `[MCP Reinitialize] OAuth required for ${serverName} (isOAuthError: ${isOAuthError}, oauthRequired: ${oauthRequired}, isOAuthFlowInitiated: ${isOAuthFlowInitiated})`, + ); + oauthRequired = true; + // Don't return error - continue so frontend can handle OAuth + } else { + logger.error( + `[MCP Reinitialize] Error initializing MCP server ${serverName} for user:`, + err, + ); + return res.status(500).json({ error: 'Failed to reinitialize MCP server for user' }); } } - // Add the new tools from this server - const tools = await userConnection.fetchTools(); - for (const tool of tools) { - const name = `${tool.name}${Constants.mcp_delimiter}${serverName}`; - userTools[name] = { - type: 'function', - ['function']: { - name, - description: tool.description, - parameters: tool.inputSchema, - }, - }; + // Only fetch and cache tools if we successfully connected (no OAuth required) + if (userConnection && !oauthRequired) { + const userTools = (await getCachedTools({ userId: user.id })) || {}; + + // Remove any old tools from this server in the user's cache + const mcpDelimiter = Constants.mcp_delimiter; + for (const key of Object.keys(userTools)) { + if (key.endsWith(`${mcpDelimiter}${serverName}`)) { + delete userTools[key]; + } + } + + // Add the new tools from this server + const tools = await userConnection.fetchTools(); + for (const tool of tools) { + const name = `${tool.name}${Constants.mcp_delimiter}${serverName}`; + userTools[name] = { + type: 'function', + ['function']: { + name, + description: tool.description, + parameters: tool.inputSchema, + }, + }; + } + + // Save the updated user tool cache + await setCachedTools(userTools, { userId: user.id }); } - // Save the updated user tool cache - await setCachedTools(userTools, { userId: user.id }); + logger.debug( + `[MCP Reinitialize] Sending response for ${serverName} - oauthRequired: ${oauthRequired}, oauthUrl: ${oauthUrl ? 'present' : 'null'}`, + ); res.json({ success: true, - message: `MCP server '${serverName}' reinitialized successfully`, + message: oauthRequired + ? `MCP server '${serverName}' ready for OAuth authentication` + : `MCP server '${serverName}' reinitialized successfully`, serverName, + oauthRequired, + oauthUrl, }); } catch (error) { logger.error('[MCP Reinitialize] Unexpected error', error); @@ -307,4 +465,163 @@ router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => { } }); +/** + * Get connection status for all MCP servers + * This endpoint returns the actual connection status from MCPManager without disconnecting idle connections + */ +router.get('/connection/status', requireJwtAuth, async (req, res) => { + try { + const user = req.user; + + if (!user?.id) { + return res.status(401).json({ error: 'User not authenticated' }); + } + + const mcpManager = getMCPManager(user.id); + const connectionStatus = {}; + + const printConfig = false; + const config = await loadCustomConfig(printConfig); + const mcpConfig = config?.mcpServers; + + const appConnections = mcpManager.getAllConnections() || new Map(); + const userConnections = mcpManager.getUserConnections(user.id) || new Map(); + const oauthServers = mcpManager.getOAuthServers() || new Set(); + + if (!mcpConfig) { + return res.status(404).json({ error: 'MCP config not found' }); + } + + // Get flow manager to check for active/timed-out OAuth flows + const flowsCache = getLogStores(CacheKeys.FLOWS); + const flowManager = getFlowStateManager(flowsCache); + + for (const [serverName] of Object.entries(mcpConfig)) { + const getConnectionState = (serverName) => + appConnections.get(serverName)?.connectionState ?? + userConnections.get(serverName)?.connectionState ?? + 'disconnected'; + + const baseConnectionState = getConnectionState(serverName); + + let hasActiveOAuthFlow = false; + let hasFailedOAuthFlow = false; + + if (baseConnectionState === 'disconnected' && oauthServers.has(serverName)) { + try { + // Check for user-specific OAuth flows + const flowId = MCPOAuthHandler.generateFlowId(user.id, serverName); + const flowState = await flowManager.getFlowState(flowId, 'mcp_oauth'); + if (flowState) { + // Check if flow failed or timed out + const flowAge = Date.now() - flowState.createdAt; + const flowTTL = flowState.ttl || 180000; // Default 3 minutes + + if (flowState.status === 'FAILED' || flowAge > flowTTL) { + hasFailedOAuthFlow = true; + logger.debug(`[MCP Connection Status] Found failed OAuth flow for ${serverName}`, { + flowId, + status: flowState.status, + flowAge, + flowTTL, + timedOut: flowAge > flowTTL, + }); + } else if (flowState.status === 'PENDING') { + hasActiveOAuthFlow = true; + logger.debug(`[MCP Connection Status] Found active OAuth flow for ${serverName}`, { + flowId, + flowAge, + flowTTL, + }); + } + } + } catch (error) { + logger.error( + `[MCP Connection Status] Error checking OAuth flows for ${serverName}:`, + error, + ); + } + } + + // Determine the final connection state + let finalConnectionState = baseConnectionState; + if (hasFailedOAuthFlow) { + finalConnectionState = 'error'; // Report as error if OAuth failed + } else if (hasActiveOAuthFlow && baseConnectionState === 'disconnected') { + finalConnectionState = 'connecting'; // Still waiting for OAuth + } + + connectionStatus[serverName] = { + requiresOAuth: oauthServers.has(serverName), + connectionState: finalConnectionState, + }; + } + + res.json({ + success: true, + connectionStatus, + }); + } catch (error) { + logger.error('[MCP Connection Status] Failed to get connection status', error); + res.status(500).json({ error: 'Failed to get connection status' }); + } +}); + +/** + * Check which authentication values exist for a specific MCP server + * This endpoint returns only boolean flags indicating if values are set, not the actual values + */ +router.get('/:serverName/auth-values', requireJwtAuth, async (req, res) => { + try { + const { serverName } = req.params; + const user = req.user; + + if (!user?.id) { + return res.status(401).json({ error: 'User not authenticated' }); + } + + 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`, + }); + } + + const serverConfig = config.mcpServers[serverName]; + const pluginKey = `${Constants.mcp_prefix}${serverName}`; + const authValueFlags = {}; + + // Check existence of saved values for each custom user variable (don't fetch actual values) + if (serverConfig.customUserVars && typeof serverConfig.customUserVars === 'object') { + for (const varName of Object.keys(serverConfig.customUserVars)) { + try { + const value = await getUserPluginAuthValue(user.id, varName, false, pluginKey); + // Only store boolean flag indicating if value exists + authValueFlags[varName] = !!(value && value.length > 0); + } catch (err) { + logger.error( + `[MCP Auth Value Flags] Error checking ${varName} for user ${user.id}:`, + err, + ); + // Default to false if we can't check + authValueFlags[varName] = false; + } + } + } + + res.json({ + success: true, + serverName, + authValueFlags, + }); + } catch (error) { + logger.error( + `[MCP Auth Value Flags] Failed to check auth value flags for ${req.params.serverName}`, + error, + ); + res.status(500).json({ error: 'Failed to check auth value flags' }); + } +}); + module.exports = router; diff --git a/api/server/services/Config/loadCustomConfig.js b/api/server/services/Config/loadCustomConfig.js index 393281daf2..2ff7d5373f 100644 --- a/api/server/services/Config/loadCustomConfig.js +++ b/api/server/services/Config/loadCustomConfig.js @@ -25,7 +25,7 @@ let i = 0; * @function loadCustomConfig * @returns {Promise} A promise that resolves to null or the custom config object. * */ -async function loadCustomConfig() { +async function loadCustomConfig(printConfig = true) { // Use CONFIG_PATH if set, otherwise fallback to defaultConfigPath const configPath = process.env.CONFIG_PATH || defaultConfigPath; @@ -108,9 +108,11 @@ https://www.librechat.ai/docs/configuration/stt_tts`); return null; } else { - logger.info('Custom config file loaded:'); - logger.info(JSON.stringify(customConfig, null, 2)); - logger.debug('Custom config:', customConfig); + if (printConfig) { + logger.info('Custom config file loaded:'); + logger.info(JSON.stringify(customConfig, null, 2)); + logger.debug('Custom config:', customConfig); + } } (customConfig.endpoints?.custom ?? []) diff --git a/api/server/services/PluginService.js b/api/server/services/PluginService.js index af42e0471c..10422d49d2 100644 --- a/api/server/services/PluginService.js +++ b/api/server/services/PluginService.js @@ -8,6 +8,7 @@ const { findOnePluginAuth, updatePluginAuth, deletePluginAuth } = require('~/mod * @param {string} userId - The unique identifier of the user for whom the plugin authentication value is to be retrieved. * @param {string} authField - The specific authentication field (e.g., 'API_KEY', 'URL') whose value is to be retrieved and decrypted. * @param {boolean} throwError - Whether to throw an error if the authentication value does not exist. Defaults to `true`. + * @param {string} [pluginKey] - Optional plugin key to make the lookup more specific to a particular plugin. * @returns {Promise} A promise that resolves to the decrypted authentication value if found, or `null` if no such authentication value exists for the given user and field. * * The function throws an error if it encounters any issue during the retrieval or decryption process, or if the authentication value does not exist. @@ -20,14 +21,28 @@ const { findOnePluginAuth, updatePluginAuth, deletePluginAuth } = require('~/mod * console.error(err); * }); * + * @example + * // To get the decrypted value of the 'API_KEY' field for a specific plugin: + * getUserPluginAuthValue('12345', 'API_KEY', true, 'mcp-server-name').then(value => { + * console.log(value); + * }).catch(err => { + * console.error(err); + * }); + * * @throws {Error} Throws an error if there's an issue during the retrieval or decryption process, or if the authentication value does not exist. * @async */ -const getUserPluginAuthValue = async (userId, authField, throwError = true) => { +const getUserPluginAuthValue = async (userId, authField, throwError = true, pluginKey) => { try { - const pluginAuth = await findOnePluginAuth({ userId, authField }); + const searchParams = { userId, authField }; + if (pluginKey) { + searchParams.pluginKey = pluginKey; + } + + const pluginAuth = await findOnePluginAuth(searchParams); if (!pluginAuth) { - throw new Error(`No plugin auth ${authField} found for user ${userId}`); + const pluginInfo = pluginKey ? ` for plugin ${pluginKey}` : ''; + throw new Error(`No plugin auth ${authField} found for user ${userId}${pluginInfo}`); } const decryptedValue = await decrypt(pluginAuth.value); diff --git a/client/src/components/Chat/Input/MCPSelect.tsx b/client/src/components/Chat/Input/MCPSelect.tsx index 0a03decd53..fd38e73c13 100644 --- a/client/src/components/Chat/Input/MCPSelect.tsx +++ b/client/src/components/Chat/Input/MCPSelect.tsx @@ -1,32 +1,45 @@ -import React, { memo, useCallback, useState } from 'react'; -import { SettingsIcon } from 'lucide-react'; -import { Constants } from 'librechat-data-provider'; -import { useUpdateUserPluginsMutation } from 'librechat-data-provider/react-query'; +import { useQueryClient } from '@tanstack/react-query'; +import { Constants, QueryKeys } from 'librechat-data-provider'; import type { TUpdateUserPlugins, TPlugin } from 'librechat-data-provider'; -import MCPConfigDialog, { type ConfigFieldDetail } from '~/components/ui/MCPConfigDialog'; +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 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'; -const getBaseMCPPluginKey = (fullPluginKey: string): string => { - const parts = fullPluginKey.split(Constants.mcp_delimiter); - return Constants.mcp_prefix + parts[parts.length - 1]; -}; - function MCPSelect() { const localize = useLocalize(); const { showToast } = useToastContext(); const { mcpSelect, startupConfig } = useBadgeRowContext(); - const { mcpValues, setMCPValues, mcpServerNames, mcpToolDetails, isPinned } = mcpSelect; + const { mcpValues, setMCPValues, mcpToolDetails, isPinned } = mcpSelect; + + // Get all configured MCP servers from config + const configuredServers = useMemo(() => { + return Object.keys(startupConfig?.mcpServers || {}); + }, [startupConfig?.mcpServers]); const [isConfigModalOpen, setIsConfigModalOpen] = useState(false); const [selectedToolForConfig, setSelectedToolForConfig] = useState(null); + const previousFocusRef = useRef(null); + + const queryClient = useQueryClient(); const updateUserPluginsMutation = useUpdateUserPluginsMutation({ - onSuccess: () => { - setIsConfigModalOpen(false); + 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); @@ -37,6 +50,61 @@ function MCPSelect() { }, }); + // 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 renderSelectedValues = useCallback( (values: string[], placeholder?: string) => { if (values.length === 0) { @@ -53,10 +121,12 @@ function MCPSelect() { const handleConfigSave = useCallback( (targetName: string, authData: Record) => { if (selectedToolForConfig && selectedToolForConfig.name === targetName) { - const basePluginKey = getBaseMCPPluginKey(selectedToolForConfig.pluginKey); - + // 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: basePluginKey, + pluginKey: `${Constants.mcp_prefix}${targetName}`, action: 'install', auth: authData, }; @@ -69,54 +139,165 @@ function MCPSelect() { const handleConfigRevoke = useCallback( (targetName: string) => { if (selectedToolForConfig && selectedToolForConfig.name === targetName) { - const basePluginKey = getBaseMCPPluginKey(selectedToolForConfig.pluginKey); - + // Use the pluginKey directly since it's already in the correct format const payload: TUpdateUserPlugins = { - pluginKey: basePluginKey, + 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], + [selectedToolForConfig, updateUserPluginsMutation, mcpValues, setMCPValues], + ); + + const handleSave = useCallback( + (authData: Record) => { + 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 hasAuthConfig = tool?.authConfig && tool.authConfig.length > 0; + 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); + }; // Common wrapper for the main content (check mark + text) // Ensures Check & Text are adjacent and the group takes available space. const mainContentWrapper = ( -
{defaultContent}
+ ); - if (tool && hasAuthConfig) { + // Check if this server has customUserVars to configure + const hasCustomUserVars = + serverConfig?.customUserVars && Object.keys(serverConfig.customUserVars).length > 0; + + const statusIcon = ( + + ); + + if (statusIcon) { return (
{mainContentWrapper} - +
{statusIcon}
); } - // For items without a settings icon, return the consistently wrapped main content. + return mainContentWrapper; }, - [mcpToolDetails, setSelectedToolForConfig, setIsConfigModalOpen], + [ + isInitializing, + isCancellable, + mcpToolDetails, + cancelOAuthFlow, + connectionStatus, + startupConfig?.mcpServers, + ], ); // Don't render if no servers are selected and not pinned @@ -124,7 +305,8 @@ function MCPSelect() { return null; } - if (!mcpToolDetails || mcpToolDetails.length === 0) { + // Don't render if no MCP servers are configured + if (!configuredServers || configuredServers.length === 0) { return null; } @@ -133,9 +315,9 @@ function MCPSelect() { return ( <> {selectedToolForConfig && ( { const schema: Record = {}; if (selectedToolForConfig?.authConfig) { @@ -173,16 +356,8 @@ function MCPSelect() { } return initial; })()} - onSave={(authData) => { - if (selectedToolForConfig) { - handleConfigSave(selectedToolForConfig.name, authData); - } - }} - onRevoke={() => { - if (selectedToolForConfig) { - handleConfigRevoke(selectedToolForConfig.name); - } - }} + onSave={handleSave} + onRevoke={handleRevoke} isSubmitting={updateUserPluginsMutation.isLoading} /> )} diff --git a/client/src/components/SidePanel/MCP/MCPPanel.tsx b/client/src/components/SidePanel/MCP/MCPPanel.tsx index 0a8ca856f6..e6a5266f84 100644 --- a/client/src/components/SidePanel/MCP/MCPPanel.tsx +++ b/client/src/components/SidePanel/MCP/MCPPanel.tsx @@ -1,60 +1,40 @@ -import { Constants } from 'librechat-data-provider'; -import { ChevronLeft, RefreshCw } from 'lucide-react'; -import { useForm, Controller } from 'react-hook-form'; -import React, { useState, useCallback, useMemo, useEffect } from 'react'; -import { - useUpdateUserPluginsMutation, - useReinitializeMCPServerMutation, -} from 'librechat-data-provider/react-query'; +import { ChevronLeft } from 'lucide-react'; +import { useQueryClient } from '@tanstack/react-query'; +import React, { useState, useMemo, useCallback } from 'react'; +import { Constants, QueryKeys } from 'librechat-data-provider'; import type { TUpdateUserPlugins } from 'librechat-data-provider'; -import { Button, Input, Label } from '~/components/ui'; +import { useUpdateUserPluginsMutation } from 'librechat-data-provider/react-query'; +import ServerInitializationSection from '~/components/ui/MCP/ServerInitializationSection'; +import CustomUserVarsSection from '~/components/ui/MCP/CustomUserVarsSection'; +import { useMCPConnectionStatusQuery } from '~/data-provider/Tools/queries'; import { useGetStartupConfig } from '~/data-provider'; import MCPPanelSkeleton from './MCPPanelSkeleton'; import { useToastContext } from '~/Providers'; +import { Button } from '~/components/ui'; import { useLocalize } from '~/hooks'; -interface ServerConfigWithVars { - serverName: string; - config: { - customUserVars: Record; - }; -} - export default function MCPPanel() { const localize = useLocalize(); const { showToast } = useToastContext(); + const queryClient = useQueryClient(); const { data: startupConfig, isLoading: startupConfigLoading } = useGetStartupConfig(); + const { data: connectionStatusData } = useMCPConnectionStatusQuery(); const [selectedServerNameForEditing, setSelectedServerNameForEditing] = useState( null, ); - const [rotatingServers, setRotatingServers] = useState>(new Set()); - const reinitializeMCPMutation = useReinitializeMCPServerMutation(); - - const mcpServerDefinitions = useMemo(() => { - if (!startupConfig?.mcpServers) { - return []; - } - return Object.entries(startupConfig.mcpServers) - .filter( - ([, serverConfig]) => - serverConfig.customUserVars && Object.keys(serverConfig.customUserVars).length > 0, - ) - .map(([serverName, config]) => ({ - serverName, - iconPath: null, - config: { - ...config, - customUserVars: config.customUserVars ?? {}, - }, - })); - }, [startupConfig?.mcpServers]); const updateUserPluginsMutation = useUpdateUserPluginsMutation({ - onSuccess: () => { + 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) => { - console.error('Error updating MCP custom user variables:', error); + onError: (error: unknown) => { + console.error('Error updating MCP auth:', error); showToast({ message: localize('com_nav_mcp_vars_update_error'), status: 'error', @@ -62,28 +42,23 @@ export default function MCPPanel() { }, }); - const handleSaveServerVars = useCallback( - (serverName: string, updatedValues: Record) => { - const payload: TUpdateUserPlugins = { - pluginKey: `${Constants.mcp_prefix}${serverName}`, - action: 'install', // 'install' action is used to set/update credentials/variables - auth: updatedValues, - }; - updateUserPluginsMutation.mutate(payload); - }, - [updateUserPluginsMutation], - ); + const mcpServerDefinitions = useMemo(() => { + if (!startupConfig?.mcpServers) { + return []; + } + return Object.entries(startupConfig.mcpServers).map(([serverName, config]) => ({ + serverName, + iconPath: null, + config: { + ...config, + customUserVars: config.customUserVars ?? {}, + }, + })); + }, [startupConfig?.mcpServers]); - const handleRevokeServerVars = useCallback( - (serverName: string) => { - const payload: TUpdateUserPlugins = { - pluginKey: `${Constants.mcp_prefix}${serverName}`, - action: 'uninstall', // 'uninstall' action clears the variables - auth: {}, // Empty auth for uninstall - }; - updateUserPluginsMutation.mutate(payload); - }, - [updateUserPluginsMutation], + const connectionStatus = useMemo( + () => connectionStatusData?.connectionStatus || {}, + [connectionStatusData?.connectionStatus], ); const handleServerClickToEdit = (serverName: string) => { @@ -94,30 +69,31 @@ export default function MCPPanel() { setSelectedServerNameForEditing(null); }; - const handleReinitializeServer = useCallback( - async (serverName: string) => { - setRotatingServers((prev) => new Set(prev).add(serverName)); - try { - await reinitializeMCPMutation.mutateAsync(serverName); - showToast({ - message: `MCP server '${serverName}' reinitialized successfully`, - status: 'success', - }); - } catch (error) { - console.error('Error reinitializing MCP server:', error); - showToast({ - message: 'Failed to reinitialize MCP server', - status: 'error', - }); - } finally { - setRotatingServers((prev) => { - const next = new Set(prev); - next.delete(serverName); - return next; - }); - } + const handleConfigSave = useCallback( + (targetName: string, authData: Record) => { + console.log( + `[MCP Panel] Saving config for ${targetName}, pluginKey: ${`${Constants.mcp_prefix}${targetName}`}`, + ); + const payload: TUpdateUserPlugins = { + pluginKey: `${Constants.mcp_prefix}${targetName}`, + action: 'install', + auth: authData, + }; + updateUserPluginsMutation.mutate(payload); }, - [showToast, reinitializeMCPMutation], + [updateUserPluginsMutation], + ); + + const handleConfigRevoke = useCallback( + (targetName: string) => { + const payload: TUpdateUserPlugins = { + pluginKey: `${Constants.mcp_prefix}${targetName}`, + action: 'uninstall', + auth: {}, + }; + updateUserPluginsMutation.mutate(payload); + }, + [updateUserPluginsMutation], ); if (startupConfigLoading) { @@ -148,6 +124,8 @@ export default function MCPPanel() { ); } + const serverStatus = connectionStatus[selectedServerNameForEditing]; + return (
+

{localize('com_sidepanel_mcp_variables_for', { '0': serverBeingEdited.serverName })}

- + +
+ + {/* Custom User Variables Section */} + { + if (selectedServerNameForEditing) { + handleConfigSave(selectedServerNameForEditing, authData); + } + }} + onRevoke={() => { + if (selectedServerNameForEditing) { + handleConfigRevoke(selectedServerNameForEditing); + } + }} isSubmitting={updateUserPluginsMutation.isLoading} /> @@ -174,124 +172,37 @@ export default function MCPPanel() { return (
- {mcpServerDefinitions.map((server) => ( -
- - -
- ))} + {mcpServerDefinitions.map((server) => { + const serverStatus = connectionStatus[server.serverName]; + const isConnected = serverStatus?.connectionState === 'connected'; + + return ( +
+ +
+ ); + })}
); } } - -// Inner component for the form - remains the same -interface MCPVariableEditorProps { - server: ServerConfigWithVars; - onSave: (serverName: string, updatedValues: Record) => void; - onRevoke: (serverName: string) => void; - isSubmitting: boolean; -} - -function MCPVariableEditor({ server, onSave, onRevoke, isSubmitting }: MCPVariableEditorProps) { - const localize = useLocalize(); - - const { - control, - handleSubmit, - reset, - formState: { errors, isDirty }, - } = useForm>({ - defaultValues: {}, // Initialize empty, will be reset by useEffect - }); - - useEffect(() => { - // Always initialize with empty strings based on the schema - const initialFormValues = Object.keys(server.config.customUserVars).reduce( - (acc, key) => { - acc[key] = ''; - return acc; - }, - {} as Record, - ); - reset(initialFormValues); - }, [reset, server.config.customUserVars]); - - const onFormSubmit = (data: Record) => { - onSave(server.serverName, data); - }; - - const handleRevokeClick = () => { - onRevoke(server.serverName); - }; - - return ( -
- {Object.entries(server.config.customUserVars).map(([key, details]) => ( -
- - ( - - )} - /> - {details.description && ( -

- )} - {errors[key] &&

{errors[key]?.message}

} -
- ))} -
- {Object.keys(server.config.customUserVars).length > 0 && ( - - )} - -
-
- ); -} diff --git a/client/src/components/ui/MCP/CustomUserVarsSection.tsx b/client/src/components/ui/MCP/CustomUserVarsSection.tsx new file mode 100644 index 0000000000..865ef91ba3 --- /dev/null +++ b/client/src/components/ui/MCP/CustomUserVarsSection.tsx @@ -0,0 +1,161 @@ +import React, { useMemo } from 'react'; +import { useForm, Controller } from 'react-hook-form'; +import { Input, Label, Button } from '~/components/ui'; +import { useLocalize } from '~/hooks'; +import { useMCPAuthValuesQuery } from '~/data-provider/Tools/queries'; + +export interface CustomUserVarConfig { + title: string; + description?: string; +} + +interface CustomUserVarsSectionProps { + serverName: string; + fields: Record; + onSave: (authData: Record) => void; + onRevoke: () => void; + isSubmitting?: boolean; +} + +interface AuthFieldProps { + name: string; + config: CustomUserVarConfig; + hasValue: boolean; + control: any; + errors: any; +} + +function AuthField({ name, config, hasValue, control, errors }: AuthFieldProps) { + const localize = useLocalize(); + + return ( +
+
+ + {hasValue ? ( +
+
+ {localize('com_ui_set')} +
+ ) : ( +
+
+ {localize('com_ui_unset')} +
+ )} +
+ ( + + )} + /> + {config.description && ( +

+ )} + {errors[name] &&

{errors[name]?.message}

} +
+ ); +} + +export default function CustomUserVarsSection({ + serverName, + fields, + onSave, + onRevoke, + isSubmitting = false, +}: CustomUserVarsSectionProps) { + const localize = useLocalize(); + + // Fetch auth value flags for the server + const { data: authValuesData } = useMCPAuthValuesQuery(serverName, { + enabled: !!serverName, + }); + + const { + control, + handleSubmit, + reset, + formState: { errors }, + } = useForm>({ + defaultValues: useMemo(() => { + const initial: Record = {}; + Object.keys(fields).forEach((key) => { + initial[key] = ''; + }); + return initial; + }, [fields]), + }); + + const onFormSubmit = (data: Record) => { + onSave(data); + }; + + const handleRevokeClick = () => { + onRevoke(); + // Reset form after revoke + reset(); + }; + + // Don't render if no fields to configure + if (!fields || Object.keys(fields).length === 0) { + return null; + } + + return ( +
+
+ {Object.entries(fields).map(([key, config]) => { + const hasValue = authValuesData?.authValueFlags?.[key] || false; + + return ( + + ); + })} + + +
+ + +
+
+ ); +} diff --git a/client/src/components/ui/MCP/MCPConfigDialog.tsx b/client/src/components/ui/MCP/MCPConfigDialog.tsx new file mode 100644 index 0000000000..fc0550453e --- /dev/null +++ b/client/src/components/ui/MCP/MCPConfigDialog.tsx @@ -0,0 +1,139 @@ +import React from 'react'; +import { Loader2, KeyRound, PlugZap, AlertTriangle } from 'lucide-react'; +import { MCPServerStatus } from 'librechat-data-provider/dist/types/types/queries'; +import { + OGDialog, + OGDialogContent, + OGDialogHeader, + OGDialogTitle, + OGDialogDescription, +} from '~/components/ui/OriginalDialog'; +import CustomUserVarsSection from './CustomUserVarsSection'; +import ServerInitializationSection from './ServerInitializationSection'; +import { useLocalize } from '~/hooks'; + +export interface ConfigFieldDetail { + title: string; + description: string; +} + +interface MCPConfigDialogProps { + isOpen: boolean; + onOpenChange: (isOpen: boolean) => void; + fieldsSchema: Record; + initialValues: Record; + onSave: (updatedValues: Record) => void; + isSubmitting?: boolean; + onRevoke?: () => void; + serverName: string; + serverStatus?: MCPServerStatus; +} + +export default function MCPConfigDialog({ + isOpen, + onOpenChange, + fieldsSchema, + onSave, + isSubmitting = false, + onRevoke, + serverName, + serverStatus, +}: MCPConfigDialogProps) { + const localize = useLocalize(); + + const hasFields = Object.keys(fieldsSchema).length > 0; + const dialogTitle = hasFields + ? localize('com_ui_configure_mcp_variables_for', { 0: serverName }) + : `${serverName} MCP Server`; + const dialogDescription = hasFields + ? localize('com_ui_mcp_dialog_desc') + : `Manage connection and settings for the ${serverName} MCP server.`; + + // Helper function to render status badge based on connection state + const renderStatusBadge = () => { + if (!serverStatus) { + return null; + } + + const { connectionState, requiresOAuth } = serverStatus; + + if (connectionState === 'connecting') { + return ( +
+ + {localize('com_ui_connecting')} +
+ ); + } + + if (connectionState === 'disconnected') { + if (requiresOAuth) { + return ( +
+ + {localize('com_ui_oauth')} +
+ ); + } else { + return ( +
+ + {localize('com_ui_offline')} +
+ ); + } + } + + if (connectionState === 'error') { + return ( +
+ + {localize('com_ui_error')} +
+ ); + } + + if (connectionState === 'connected') { + return ( +
+
+ {localize('com_ui_active')} +
+ ); + } + + return null; + }; + + return ( + + + +
+ {dialogTitle} + {renderStatusBadge()} +
+ {dialogDescription} + + + {/* Content */} +
+ {/* Custom User Variables Section */} + {})} + isSubmitting={isSubmitting} + /> +
+ + {/* Server Initialization Section */} + + + + ); +} diff --git a/client/src/components/ui/MCP/MCPServerStatusIcon.tsx b/client/src/components/ui/MCP/MCPServerStatusIcon.tsx new file mode 100644 index 0000000000..69f328a0c8 --- /dev/null +++ b/client/src/components/ui/MCP/MCPServerStatusIcon.tsx @@ -0,0 +1,190 @@ +import React from 'react'; +import { SettingsIcon, AlertTriangle, Loader2, KeyRound, PlugZap, X } from 'lucide-react'; +import type { MCPServerStatus, TPlugin } from 'librechat-data-provider'; +import { useLocalize } from '~/hooks'; + +let localize: ReturnType; + +interface StatusIconProps { + serverName: string; + onConfigClick: (e: React.MouseEvent) => void; +} + +interface InitializingStatusProps extends StatusIconProps { + onCancel: (e: React.MouseEvent) => void; + canCancel: boolean; +} + +interface MCPServerStatusIconProps { + serverName: string; + serverStatus?: MCPServerStatus; + tool?: TPlugin; + onConfigClick: (e: React.MouseEvent) => void; + isInitializing: boolean; + canCancel: boolean; + onCancel: (e: React.MouseEvent) => void; + hasCustomUserVars?: boolean; +} + +/** + * Renders the appropriate status icon for an MCP server based on its state + */ +export default function MCPServerStatusIcon({ + serverName, + serverStatus, + tool, + onConfigClick, + isInitializing, + canCancel, + onCancel, + hasCustomUserVars = false, +}: MCPServerStatusIconProps) { + localize = useLocalize(); + if (isInitializing) { + return ( + + ); + } + + if (!serverStatus) { + return null; + } + + const { connectionState, requiresOAuth } = serverStatus; + + if (connectionState === 'connecting') { + return ; + } + + if (connectionState === 'disconnected') { + if (requiresOAuth) { + return ; + } + return ; + } + + if (connectionState === 'error') { + return ; + } + + if (connectionState === 'connected') { + // Only show config button if there are customUserVars to configure + if (hasCustomUserVars) { + const isAuthenticated = tool?.authenticated || requiresOAuth; + return ( + + ); + } + return null; // No config button for connected servers without customUserVars + } + + return null; +} + +function InitializingStatusIcon({ serverName, onCancel, canCancel }: InitializingStatusProps) { + if (canCancel) { + return ( + + ); + } + + return ( +
+ +
+ ); +} + +function ConnectingStatusIcon({ serverName }: StatusIconProps) { + return ( +
+ +
+ ); +} + +function DisconnectedOAuthStatusIcon({ serverName, onConfigClick }: StatusIconProps) { + return ( + + ); +} + +function DisconnectedStatusIcon({ serverName, onConfigClick }: StatusIconProps) { + return ( + + ); +} + +function ErrorStatusIcon({ serverName, onConfigClick }: StatusIconProps) { + return ( + + ); +} + +interface AuthenticatedStatusProps extends StatusIconProps { + isAuthenticated: boolean; +} + +function AuthenticatedStatusIcon({ + serverName, + onConfigClick, + isAuthenticated, +}: AuthenticatedStatusProps) { + return ( + + ); +} diff --git a/client/src/components/ui/MCP/ServerInitializationSection.tsx b/client/src/components/ui/MCP/ServerInitializationSection.tsx new file mode 100644 index 0000000000..dd2588c2f2 --- /dev/null +++ b/client/src/components/ui/MCP/ServerInitializationSection.tsx @@ -0,0 +1,131 @@ +import { RefreshCw, Link } from 'lucide-react'; +import React, { useState, useCallback } from 'react'; +import { useMCPServerInitialization } from '~/hooks/MCP/useMCPServerInitialization'; +import { Button } from '~/components/ui'; +import { useLocalize } from '~/hooks'; + +interface ServerInitializationSectionProps { + serverName: string; + requiresOAuth: boolean; +} + +export default function ServerInitializationSection({ + serverName, + requiresOAuth, +}: ServerInitializationSectionProps) { + const localize = useLocalize(); + + const [oauthUrl, setOauthUrl] = useState(null); + + // Use the shared initialization hook + const { initializeServer, isLoading, connectionStatus, cancelOAuthFlow, isCancellable } = + useMCPServerInitialization({ + onOAuthStarted: (name, url) => { + // Store the OAuth URL locally for display + setOauthUrl(url); + }, + onSuccess: () => { + // Clear OAuth URL on success + setOauthUrl(null); + }, + }); + + const serverStatus = connectionStatus[serverName]; + const isConnected = serverStatus?.connectionState === 'connected'; + const canCancel = isCancellable(serverName); + + const handleInitializeClick = useCallback(() => { + setOauthUrl(null); + initializeServer(serverName); + }, [initializeServer, serverName]); + + const handleCancelClick = useCallback(() => { + setOauthUrl(null); + cancelOAuthFlow(serverName); + }, [cancelOAuthFlow, serverName]); + + // Show subtle reinitialize option if connected + if (isConnected) { + return ( +
+ +
+ ); + } + + return ( +
+
+
+ + {requiresOAuth + ? localize('com_ui_mcp_not_authenticated', { 0: serverName }) + : localize('com_ui_mcp_not_initialized', { 0: serverName })} + +
+ {/* Only show authenticate button when OAuth URL is not present */} + {!oauthUrl && ( + + )} +
+ + {/* OAuth URL display */} + {oauthUrl && ( +
+
+
+ +
+ + {localize('com_ui_auth_url')} + +
+
+ + +
+

+ {localize('com_ui_oauth_flow_desc')} +

+
+ )} +
+ ); +} diff --git a/client/src/components/ui/MCPConfigDialog.tsx b/client/src/components/ui/MCPConfigDialog.tsx deleted file mode 100644 index d1a53bd902..0000000000 --- a/client/src/components/ui/MCPConfigDialog.tsx +++ /dev/null @@ -1,122 +0,0 @@ -import React, { useEffect } from 'react'; -import { useForm, Controller } from 'react-hook-form'; -import { Input, Label, OGDialog, Button } from '~/components/ui'; -import OGDialogTemplate from '~/components/ui/OGDialogTemplate'; -import { useLocalize } from '~/hooks'; - -export interface ConfigFieldDetail { - title: string; - description: string; -} - -interface MCPConfigDialogProps { - isOpen: boolean; - onOpenChange: (isOpen: boolean) => void; - fieldsSchema: Record; - initialValues: Record; - onSave: (updatedValues: Record) => void; - isSubmitting?: boolean; - onRevoke?: () => void; - serverName: string; -} - -export default function MCPConfigDialog({ - isOpen, - onOpenChange, - fieldsSchema, - initialValues, - onSave, - isSubmitting = false, - onRevoke, - serverName, -}: MCPConfigDialogProps) { - const localize = useLocalize(); - const { - control, - handleSubmit, - reset, - formState: { errors, _ }, - } = useForm>({ - defaultValues: initialValues, - }); - - useEffect(() => { - if (isOpen) { - reset(initialValues); - } - }, [isOpen, initialValues, reset]); - - const onFormSubmit = (data: Record) => { - onSave(data); - }; - - const handleRevoke = () => { - if (onRevoke) { - onRevoke(); - } - }; - - const dialogTitle = localize('com_ui_configure_mcp_variables_for', { 0: serverName }); - const dialogDescription = localize('com_ui_mcp_dialog_desc'); - - return ( - - - {Object.entries(fieldsSchema).map(([key, details]) => ( -
- - ( - - )} - /> - {details.description && ( -

- )} - {errors[key] &&

{errors[key]?.message}

} -
- ))} - - } - selection={{ - selectHandler: handleSubmit(onFormSubmit), - selectClasses: 'bg-green-500 hover:bg-green-600 text-white', - selectText: isSubmitting ? localize('com_ui_saving') : localize('com_ui_save'), - }} - buttons={ - onRevoke && ( - - ) - } - footerClassName="flex justify-end gap-2 px-6 pb-6 pt-2" - showCancelButton={true} - /> - - ); -} diff --git a/client/src/components/ui/MultiSelect.tsx b/client/src/components/ui/MultiSelect.tsx index e0b74a4577..20e2512794 100644 --- a/client/src/components/ui/MultiSelect.tsx +++ b/client/src/components/ui/MultiSelect.tsx @@ -113,7 +113,7 @@ export default function MultiSelect({ {items.map((value) => { const defaultContent = ( <> - + {value} ); diff --git a/client/src/data-provider/Tools/queries.ts b/client/src/data-provider/Tools/queries.ts index 1aea8e9f5a..6b06f25c9f 100644 --- a/client/src/data-provider/Tools/queries.ts +++ b/client/src/data-provider/Tools/queries.ts @@ -40,3 +40,36 @@ export const useGetToolCalls = ( }, ); }; + +export const useMCPConnectionStatusQuery = ( + config?: UseQueryOptions, +): QueryObserverResult => { + return useQuery( + [QueryKeys.mcpConnectionStatus], + () => dataService.getMCPConnectionStatus(), + { + refetchOnWindowFocus: false, + refetchOnReconnect: false, + refetchOnMount: false, + staleTime: 10000, // 10 seconds + ...config, + }, + ); +}; + +export const useMCPAuthValuesQuery = ( + serverName: string, + config?: UseQueryOptions, +): QueryObserverResult => { + return useQuery( + [QueryKeys.mcpAuthValues, serverName], + () => dataService.getMCPAuthValues(serverName), + { + refetchOnWindowFocus: false, + refetchOnReconnect: false, + refetchOnMount: false, + enabled: !!serverName, + ...config, + }, + ); +}; diff --git a/client/src/hooks/MCP/index.ts b/client/src/hooks/MCP/index.ts new file mode 100644 index 0000000000..134baaaad0 --- /dev/null +++ b/client/src/hooks/MCP/index.ts @@ -0,0 +1 @@ +export { useMCPServerInitialization } from './useMCPServerInitialization'; diff --git a/client/src/hooks/MCP/useMCPServerInitialization.ts b/client/src/hooks/MCP/useMCPServerInitialization.ts new file mode 100644 index 0000000000..17db7d9089 --- /dev/null +++ b/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>(new Map()); + const [oauthStartTimes, setOauthStartTimes] = useState>(new Map()); + const [initializingServers, setInitializingServers] = useState>(new Set()); + const [cancellableServers, setCancellableServers] = useState>(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, + }; +} diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 79d6bc4dcf..c7d13cd306 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -446,6 +446,8 @@ "com_nav_maximize_chat_space": "Maximize chat space", "com_nav_mcp_vars_update_error": "Error updating MCP custom user variables: {{0}}", "com_nav_mcp_vars_updated": "MCP custom user variables updated successfully.", + "com_nav_mcp_status_connecting": "{{0}} - Connecting", + "com_nav_mcp_configure_server": "Configure {{0}}", "com_nav_modular_chat": "Enable switching Endpoints mid-conversation", "com_nav_my_files": "My Files", "com_nav_not_supported": "Not Supported", @@ -503,7 +505,6 @@ "com_sidepanel_conversation_tags": "Bookmarks", "com_sidepanel_hide_panel": "Hide Panel", "com_sidepanel_manage_files": "Manage Files", - "com_sidepanel_mcp_enter_value": "Enter value for {{0}}", "com_sidepanel_mcp_no_servers_with_vars": "No MCP servers with configurable variables.", "com_sidepanel_mcp_variables_for": "MCP Variables for {{0}}", "com_sidepanel_parameters": "Parameters", @@ -1088,5 +1089,25 @@ "com_ui_x_selected": "{{0}} selected", "com_ui_yes": "Yes", "com_ui_zoom": "Zoom", + "com_user_message": "You", + "com_ui_mcp_update_var": "Update {{0}}", + "com_ui_set": "Set", + "com_ui_unset": "Unset", + "com_ui_connecting": "Connecting", + "com_ui_offline": "Offline", + "com_ui_continue_oauth": "Continue with OAuth", + "com_ui_oauth_flow_desc": "Complete the OAuth flow in the new window, then return here.", + "com_ui_authenticate": "Authenticate", + "com_ui_mcp_initialize": "Initialize", + "com_ui_mcp_not_authenticated": "{{0}} not authenticated (OAuth Required)", + "com_ui_mcp_not_initialized": "{{0}} not initialized", + "com_ui_reinitialize": "Reinitialize", + "com_ui_mcp_initialized_success": "MCP server '{{0}}' initialized successfully", + "com_ui_mcp_authenticated_success": "MCP server '{{0}}' authenticated successfully", + "com_ui_mcp_oauth_no_url": "OAuth authentication required but no URL provided", + "com_ui_mcp_oauth_cancelled": "OAuth login cancelled for {{0}}", + "com_ui_mcp_oauth_timeout": "OAuth login timed out for {{0}}", + "com_ui_mcp_init_failed": "Failed to initialize MCP server", + "com_ui_active": "Active", "com_user_message": "You" } diff --git a/packages/api/src/mcp/manager.ts b/packages/api/src/mcp/manager.ts index a3bee8a18c..6db52ec9ec 100644 --- a/packages/api/src/mcp/manager.ts +++ b/packages/api/src/mcp/manager.ts @@ -28,6 +28,8 @@ export class MCPManager { private mcpConfigs: t.MCPServers = {}; /** Store MCP server instructions */ private serverInstructions: Map = new Map(); + /** Track servers that required OAuth at startup */ + private oauthServers: Set = new Set(); public static getInstance(): MCPManager { if (!MCPManager.instance) { @@ -167,38 +169,14 @@ export class MCPManager { } const connection = new MCPConnection(serverName, processedConfig, undefined, tokens); logger.info(`[MCP][${serverName}] Setting up OAuth event listener`); - connection.on('oauthRequired', async (data) => { + connection.on('oauthRequired', async () => { logger.debug(`[MCP][${serverName}] oauthRequired event received`); - const result = await this.handleOAuthRequired({ - ...data, - flowManager, - }); - if (result?.tokens && tokenMethods?.createToken) { - try { - connection.setOAuthTokens(result.tokens); - await MCPTokenStorage.storeTokens({ - userId: CONSTANTS.SYSTEM_USER_ID, - serverName, - tokens: result.tokens, - createToken: tokenMethods.createToken, - updateToken: tokenMethods.updateToken, - findToken: tokenMethods.findToken, - clientInfo: result.clientInfo, - }); - logger.info(`[MCP][${serverName}] OAuth tokens saved to storage`); - } catch (error) { - logger.error(`[MCP][${serverName}] Failed to save OAuth tokens to storage`, error); - } - } - // Only emit oauthHandled if we actually got tokens (OAuth succeeded) - if (result?.tokens) { - connection.emit('oauthHandled'); - } else { - // OAuth failed, emit oauthFailed to properly reject the promise - logger.warn(`[MCP][${serverName}] OAuth failed, emitting oauthFailed event`); - connection.emit('oauthFailed', new Error('OAuth authentication failed')); - } + this.oauthServers.add(serverName); + + // Skip OAuth at startup - let connection fail gracefully + logger.info(`[MCP][${serverName}] OAuth required, skipping at startup`); + connection.emit('oauthFailed', new Error('OAuth authentication skipped at startup')); }); try { const connectTimeout = processedConfig.initTimeout ?? 30000; @@ -394,6 +372,7 @@ export class MCPManager { oauthStart, oauthEnd, signal, + returnOnOAuth = false, }: { user: TUser; serverName: string; @@ -403,6 +382,7 @@ export class MCPManager { oauthStart?: (authURL: string) => Promise; oauthEnd?: () => Promise; signal?: AbortSignal; + returnOnOAuth?: boolean; }): Promise { const userId = user.id; if (!userId) { @@ -513,6 +493,48 @@ export class MCPManager { connection.on('oauthRequired', async (data) => { logger.info(`[MCP][User: ${userId}][${serverName}] oauthRequired event received`); + + // If we just want to initiate OAuth and return, handle it differently + if (returnOnOAuth) { + try { + const config = this.mcpConfigs[serverName]; + const { authorizationUrl, flowId, flowMetadata } = + await MCPOAuthHandler.initiateOAuthFlow( + serverName, + data.serverUrl || '', + userId, + config?.oauth, + ); + + // Create the flow state so the OAuth callback can find it + // We spawn this in the background without waiting for it + flowManager.createFlow(flowId, 'mcp_oauth', flowMetadata).catch(() => { + // The OAuth callback will resolve this flow, so we expect it to timeout here + // which is fine - we just need the flow state to exist + }); + + if (oauthStart) { + logger.info( + `[MCP][User: ${userId}][${serverName}] OAuth flow started, issuing authorization URL`, + ); + await oauthStart(authorizationUrl); + } + + // Emit oauthFailed to signal that connection should not proceed + // but OAuth was successfully initiated + connection?.emit('oauthFailed', new Error('OAuth flow initiated - return early')); + return; + } catch (error) { + logger.error( + `[MCP][User: ${userId}][${serverName}] Failed to initiate OAuth flow`, + error, + ); + connection?.emit('oauthFailed', new Error('OAuth initiation failed')); + return; + } + } + + // Normal OAuth handling - wait for completion const result = await this.handleOAuthRequired({ ...data, flowManager, @@ -598,7 +620,6 @@ export class MCPManager { /** Removes a specific user connection entry */ private removeUserConnection(userId: string, serverName: string): void { - // Remove connection object const userMap = this.userConnections.get(userId); if (userMap) { userMap.delete(serverName); @@ -1111,4 +1132,12 @@ ${logPrefix} Flow ID: ${newFlowId} return null; } } + public getUserConnections(userId: string) { + return this.userConnections.get(userId); + } + + /** Get servers that require OAuth */ + public getOAuthServers(): Set { + return this.oauthServers; + } } diff --git a/packages/data-provider/src/api-endpoints.ts b/packages/data-provider/src/api-endpoints.ts index 3930132ce0..38da404506 100644 --- a/packages/data-provider/src/api-endpoints.ts +++ b/packages/data-provider/src/api-endpoints.ts @@ -133,6 +133,14 @@ export const resendVerificationEmail = () => '/api/user/verify/resend'; export const plugins = () => '/api/plugins'; export const mcpReinitialize = (serverName: string) => `/api/mcp/${serverName}/reinitialize`; +export const mcpConnectionStatus = () => '/api/mcp/connection/status'; +export const mcpAuthValues = (serverName: string) => { + return `/api/mcp/${serverName}/auth-values`; +}; + +export const cancelMCPOAuth = (serverName: string) => { + return `/api/mcp/oauth/cancel/${serverName}`; +}; export const config = () => '/api/config'; diff --git a/packages/data-provider/src/data-service.ts b/packages/data-provider/src/data-service.ts index d9e8ff7d12..4356028371 100644 --- a/packages/data-provider/src/data-service.ts +++ b/packages/data-provider/src/data-service.ts @@ -145,6 +145,18 @@ export const reinitializeMCPServer = (serverName: string) => { return request.post(endpoints.mcpReinitialize(serverName)); }; +export const getMCPConnectionStatus = (): Promise => { + return request.get(endpoints.mcpConnectionStatus()); +}; + +export const getMCPAuthValues = (serverName: string): Promise => { + return request.get(endpoints.mcpAuthValues(serverName)); +}; + +export function cancelMCPOAuth(serverName: string): Promise { + return request.post(endpoints.cancelMCPOAuth(serverName), {}); +} + /* Config */ export const getStartupConfig = (): Promise< diff --git a/packages/data-provider/src/keys.ts b/packages/data-provider/src/keys.ts index ec94c0f0ff..5cb4c8195a 100644 --- a/packages/data-provider/src/keys.ts +++ b/packages/data-provider/src/keys.ts @@ -27,6 +27,8 @@ export enum QueryKeys { tools = 'tools', toolAuth = 'toolAuth', toolCalls = 'toolCalls', + mcpConnectionStatus = 'mcpConnectionStatus', + mcpAuthValues = 'mcpAuthValues', agentTools = 'agentTools', actions = 'actions', assistantDocs = 'assistantDocs', diff --git a/packages/data-provider/src/react-query/react-query-service.ts b/packages/data-provider/src/react-query/react-query-service.ts index ca7d4374fe..726ed8241c 100644 --- a/packages/data-provider/src/react-query/react-query-service.ts +++ b/packages/data-provider/src/react-query/react-query-service.ts @@ -317,7 +317,13 @@ export const useUpdateUserPluginsMutation = ( }; export const useReinitializeMCPServerMutation = (): UseMutationResult< - { success: boolean; message: string; serverName: string }, + { + success: boolean; + message: string; + serverName: string; + oauthRequired?: boolean; + oauthUrl?: string; + }, unknown, string, unknown @@ -330,6 +336,20 @@ export const useReinitializeMCPServerMutation = (): UseMutationResult< }); }; +export const useCancelMCPOAuthMutation = (): UseMutationResult< + m.CancelMCPOAuthResponse, + unknown, + string, + unknown +> => { + const queryClient = useQueryClient(); + return useMutation((serverName: string) => dataService.cancelMCPOAuth(serverName), { + onSuccess: () => { + queryClient.invalidateQueries([QueryKeys.mcpConnectionStatus]); + }, + }); +}; + export const useGetCustomConfigSpeechQuery = ( config?: UseQueryOptions, ): QueryObserverResult => { diff --git a/packages/data-provider/src/types/mutations.ts b/packages/data-provider/src/types/mutations.ts index cd6cae75c8..48e34b5d8c 100644 --- a/packages/data-provider/src/types/mutations.ts +++ b/packages/data-provider/src/types/mutations.ts @@ -368,3 +368,13 @@ export type TLogoutResponse = { }; export type LogoutOptions = MutationOptions; + +export interface AssistantInitialize { + message: string; + error?: string; +} + +export interface CancelMCPOAuthResponse { + success: boolean; + message: string; +} diff --git a/packages/data-provider/src/types/queries.ts b/packages/data-provider/src/types/queries.ts index cd792e4bda..1dd326cc99 100644 --- a/packages/data-provider/src/types/queries.ts +++ b/packages/data-provider/src/types/queries.ts @@ -124,3 +124,19 @@ export type MemoriesResponse = { tokenLimit: number | null; usagePercentage: number | null; }; + +export interface MCPServerStatus { + requiresOAuth: boolean; + connectionState: 'disconnected' | 'connecting' | 'connected' | 'error'; +} + +export interface MCPConnectionStatusResponse { + success: boolean; + connectionStatus: Record; +} + +export interface MCPAuthValuesResponse { + success: boolean; + serverName: string; + authValueFlags: Record; +}