LibreChat/api/server/services/Tools/mcp.js

const { logger } = require('@librechat/data-schemas');
const { getMissingCustomUserVars, requiresEphemeralUserConnection } = require('@librechat/api');
const { CacheKeys, Constants } = require('librechat-data-provider');
const { getMCPManager, getMCPServersRegistry, getFlowStateManager } = require('~/config');
const { findToken, createToken, updateToken, deleteTokens } = require('~/models');
const { getGraphApiToken } = require('~/server/services/GraphTokenService');
const { exchangeOboToken } = require('~/server/services/OboTokenService');
const { createOboTrustChecker } = require('~/server/services/OboPolicyService');
const { updateMCPServerTools } = require('~/server/services/Config');
const { getLogStores } = require('~/cache');

/**
 * Reinitializes an MCP server connection and discovers available tools.
 * When OAuth is required, uses discovery mode to list tools without full authentication
 * (per MCP spec, tool listing should be possible without auth).
 * @param {Object} params
 * @param {IUser} params.user - The user from the request object.
 * @param {string} params.serverName - The name of the MCP server
 * @param {boolean} params.returnOnOAuth - Whether to initiate OAuth and return, or wait for OAuth flow to finish
 * @param {AbortSignal} [params.signal] - The abort signal to handle cancellation.
 * @param {boolean} [params.forceNew]
 * @param {number} [params.connectionTimeout]
 * @param {FlowStateManager<any>} [params.flowManager]
 * @param {(authURL: string, options?: { expiresAt?: number }) => Promise<void>} [params.oauthStart]
 * @param {() => Promise<void>} [params.oauthEnd]
 * @param {import('@librechat/api').RequestBody} [params.requestBody]
 * @param {Record<string, Record<string, string>>} [params.userMCPAuthMap]
 */
async function reinitMCPServer({
  user,
  signal,
  forceNew,
  serverName,
  configServers,
  userMCPAuthMap,
  connectionTimeout,
  returnOnOAuth = true,
  oauthStart: _oauthStart,
  flowManager: _flowManager,
  serverConfig: providedConfig,
  requestBody,
  oauthEnd,
}) {
  /** @type {MCPConnection | null} */
  let connection = null;
  let serverConfig = providedConfig;
  /** @type {LCAvailableTools | null} */
  let availableTools = null;
  /** @type {ReturnType<MCPConnection['fetchTools']> | null} */
  let tools = null;
  let oauthRequired = false;
  let oauthUrl = null;
  let ephemeralServer = false;

  try {
    const registry = getMCPServersRegistry();
    serverConfig =
      serverConfig ?? (await registry.getServerConfig(serverName, user?.id, configServers));
    ephemeralServer = serverConfig ? requiresEphemeralUserConnection(serverConfig) : false;
    if (serverConfig?.inspectionFailed) {
      if (serverConfig.source === 'config') {
        logger.info(
          `[MCP Reinitialize] Config-source server ${serverName} has inspectionFailed — retry handled by config cache`,
        );
        return {
          availableTools: null,
          success: false,
          message: `MCP server '${serverName}' is still unreachable`,
          oauthRequired: false,
          serverName,
          oauthUrl: null,
          tools: null,
        };
      } else {
        logger.info(
          `[MCP Reinitialize] Server ${serverName} had failed inspection, attempting reinspection`,
        );
        try {
          const storageLocation = serverConfig.source === 'user' ? 'DB' : 'CACHE';
          await registry.reinspectServer(serverName, storageLocation, user?.id);
          logger.info(`[MCP Reinitialize] Reinspection succeeded for server: ${serverName}`);
        } catch (reinspectError) {
          logger.error(
            `[MCP Reinitialize] Reinspection failed for server ${serverName}:`,
            reinspectError,
          );
          return {
            availableTools: null,
            success: false,
            message: `MCP server '${serverName}' is still unreachable`,
            oauthRequired: false,
            serverName,
            oauthUrl: null,
            tools: null,
          };
        }
      }
    }

    const customUserVars = userMCPAuthMap?.[`${Constants.mcp_prefix}${serverName}`];

    const missingUserVars = getMissingCustomUserVars(serverConfig ?? {}, customUserVars);
    if (missingUserVars.length > 0) {
      logger.warn(
        `[MCP Reinitialize] Skipping server '${serverName}': required user-provided variable(s) not set: ${missingUserVars.join(
          ', ',
        )}. Tools will not be exposed until the user configures them.`,
      );
      return {
        availableTools: null,
        success: false,
        message: `MCP server '${serverName}' requires user-provided variable(s) [${missingUserVars.join(
          ', ',
        )}] which are not set`,
        oauthRequired: false,
        serverName,
        oauthUrl: null,
        tools: null,
      };
    }

    const flowManager = _flowManager ?? getFlowStateManager(getLogStores(CacheKeys.FLOWS));
    const mcpManager = getMCPManager();
    const tokenMethods = { findToken, updateToken, createToken, deleteTokens };

    const oauthStart =
      _oauthStart ??
      (async (authURL) => {
        logger.info(`[MCP Reinitialize] OAuth URL received for ${serverName}`);
        oauthUrl = authURL;
        oauthRequired = true;
      });

    try {
      connection = await mcpManager.getConnection({
        user,
        signal,
        forceNew,
        oauthStart,
        serverName,
        flowManager,
        tokenMethods,
        returnOnOAuth,
        oauthEnd,
        customUserVars,
        requestBody,
        connectionTimeout,
        serverConfig,
        graphTokenResolver: getGraphApiToken,
        oboTokenResolver: exchangeOboToken,
        oboTrustChecker: createOboTrustChecker(),
      });

      logger.info(`[MCP Reinitialize] Successfully established connection for ${serverName}`);
    } catch (err) {
      logger.info(`[MCP Reinitialize] getConnection threw error: ${err.message}`);
      logger.info(
        `[MCP Reinitialize] OAuth state - oauthRequired: ${oauthRequired}, oauthUrl: ${oauthUrl ? 'present' : 'null'}`,
      );

      const isOAuthError =
        err.message?.includes('OAuth') ||
        err.message?.includes('authentication') ||
        err.message?.includes('401');

      const isOAuthFlowInitiated = err.message === 'OAuth flow initiated - return early';

      if (isOAuthError || oauthRequired || isOAuthFlowInitiated) {
        logger.info(
          `[MCP Reinitialize] OAuth required for ${serverName}, attempting tool discovery without auth`,
        );
        oauthRequired = true;

        try {
          const discoveryResult = await mcpManager.discoverServerTools({
            user,
            signal,
            serverName,
            flowManager,
            tokenMethods,
            oauthStart,
            customUserVars,
            requestBody,
            connectionTimeout,
            configServers,
            graphTokenResolver: getGraphApiToken,
            oboTokenResolver: exchangeOboToken,
            oboTrustChecker: createOboTrustChecker(),
          });

          if (discoveryResult.tools && discoveryResult.tools.length > 0) {
            tools = discoveryResult.tools;
            logger.info(
              `[MCP Reinitialize] Discovered ${tools.length} tools for ${serverName} without full auth`,
            );
          }
        } catch (discoveryErr) {
          logger.debug(
            `[MCP Reinitialize] Tool discovery failed for ${serverName}: ${discoveryErr?.message ?? String(discoveryErr)}`,
          );
        }
      } else {
        logger.error(
          `[MCP Reinitialize] Error initializing MCP server ${serverName} for user:`,
          err,
        );
      }
    }

    if (connection && !oauthRequired) {
      tools = await connection.fetchTools();
    }

    if (tools && tools.length > 0) {
      availableTools = await updateMCPServerTools({
        userId: user.id,
        serverName,
        tools,
        skipCache: ephemeralServer,
      });
    }

    logger.debug(
      `[MCP Reinitialize] Sending response for ${serverName} - oauthRequired: ${oauthRequired}, oauthUrl: ${oauthUrl ? 'present' : 'null'}`,
    );

    const getResponseMessage = () => {
      if (oauthRequired && tools && tools.length > 0) {
        return `MCP server '${serverName}' tools discovered, OAuth required for execution`;
      }
      if (oauthRequired) {
        return `MCP server '${serverName}' ready for OAuth authentication`;
      }
      if (connection) {
        return `MCP server '${serverName}' reinitialized successfully`;
      }
      return `Failed to reinitialize MCP server '${serverName}'`;
    };

    const result = {
      availableTools,
      success: Boolean(
        (connection && !oauthRequired) ||
          (oauthRequired && oauthUrl) ||
          (tools && tools.length > 0),
      ),
      message: getResponseMessage(),
      oauthRequired,
      serverName,
      oauthUrl,
      tools,
    };

    logger.debug(`[MCP Reinitialize] Response for ${serverName}:`, {
      success: result.success,
      oauthRequired: result.oauthRequired,
      oauthUrl: result.oauthUrl ? 'present' : null,
      toolsCount: tools?.length ?? 0,
    });

    return result;
  } catch (error) {
    logger.error(
      '[MCP Reinitialize] Error loading MCP Tools, servers may still be initializing:',
      error,
    );
  } finally {
    if (connection && ephemeralServer) {
      try {
        await connection.disconnect();
      } catch (error) {
        logger.warn(
          `[MCP Reinitialize] Failed to disconnect ephemeral server ${serverName}`,
          error,
        );
      }
    }
  }
}

module.exports = {
  reinitMCPServer,
};