🪐 feat: MCP OAuth 2.0 Discovery Support (#7924)

* chore: Update @modelcontextprotocol/sdk to version 1.12.3 in package.json and package-lock.json - Bump version of @modelcontextprotocol/sdk to 1.12.3 to incorporate recent updates. - Update dependencies for ajv and cross-spawn to their latest versions. - Add ajv as a new dependency in the sdk module. - Include json-schema-traverse as a new dependency in the sdk module. * feat: @librechat/auth * feat: Add crypto module exports to auth package - Introduced a new crypto module by creating index.ts in the crypto directory. - Updated the main index.ts of the auth package to export from the new crypto module. * feat: Update package dependencies and build scripts for auth package - Added @librechat/auth as a dependency in package.json and package-lock.json. - Updated build scripts to include the auth package in both frontend and bun build processes. - Removed unused mongoose and openid-client dependencies from package-lock.json for cleaner dependency management. * refactor: Migrate crypto utility functions to @librechat/auth - Replaced local crypto utility imports with the new @librechat/auth package across multiple files. - Removed the obsolete crypto.js file and its exports. - Updated relevant services and models to utilize the new encryption and decryption methods from @librechat/auth. * feat: Enhance OAuth token handling and update dependencies in auth package * chore: Remove Token model and TokenService due to restructuring of OAuth handling - Deleted the Token.js model and TokenService.js, which were responsible for managing OAuth tokens. - This change is part of a broader refactor to streamline OAuth token management and improve code organization. * refactor: imports from '@librechat/auth' to '@librechat/api' and add OAuth token handling functionality * refactor: Simplify logger usage in MCP and FlowStateManager classes * chore: fix imports * feat: Add OAuth configuration schema to MCP with token exchange method support * feat: FIRST PASS Implement MCP OAuth flow with token management and error handling - Added a new route for handling OAuth callbacks and token retrieval. - Integrated OAuth token storage and retrieval mechanisms. - Enhanced MCP connection to support automatic OAuth flow initiation on 401 errors. - Implemented dynamic client registration and metadata discovery for OAuth. - Updated MCPManager to manage OAuth tokens and handle authentication requirements. - Introduced comprehensive logging for OAuth processes and error handling. * refactor: Update MCPConnection and MCPManager to utilize new URL handling - Added a `url` property to MCPConnection for better URL management. - Refactored MCPManager to use the new `url` property instead of a deprecated method for OAuth handling. - Changed logging from info to debug level for flow manager and token methods initialization. - Improved comments for clarity on existing tokens and OAuth event listener setup. * refactor: Improve connection timeout error messages in MCPConnection and MCPManager and use initTimeout for connection - Updated the connection timeout error messages to include the duration of the timeout. - Introduced a configurable `connectTimeout` variable in both MCPConnection and MCPManager for better flexibility. * chore: cleanup MCP OAuth Token exchange handling; fix: erroneous use of flowsCache and remove verbose logs * refactor: Update MCPManager and MCPTokenStorage to use TokenMethods for token management - Removed direct token storage handling in MCPManager and replaced it with TokenMethods for better abstraction. - Refactored MCPTokenStorage methods to accept parameters for token operations, enhancing flexibility and readability. - Improved logging messages related to token persistence and retrieval processes. * refactor: Update MCP OAuth handling to use static methods and improve flow management - Refactored MCPOAuthHandler to utilize static methods for initiating and completing OAuth flows, enhancing clarity and reducing instance dependencies. - Updated MCPManager to pass flowManager explicitly to OAuth handling methods, improving flexibility in flow state management. - Enhanced comments and logging for better understanding of OAuth processes and flow state retrieval. * refactor: Integrate token methods into createMCPTool for enhanced token management * refactor: Change logging from info to debug level in MCPOAuthHandler for improved log management * chore: clean up logging * feat: first pass, auth URL from MCP OAuth flow * chore: Improve logging format for OAuth authentication URL display * chore: cleanup mcp manager comments * feat: add connection reconnection logic in MCPManager * refactor: reorganize token storage handling in MCP - Moved token storage logic from MCPManager to a new MCPTokenStorage class for better separation of concerns. - Updated imports to reflect the new token storage structure. - Enhanced methods for storing, retrieving, updating, and deleting OAuth tokens, improving overall token management. * chore: update comment for SYSTEM_USER_ID in MCPManager for clarity * feat: implement refresh token functionality in MCP - Added refresh token handling in MCPManager to support token renewal for both app-level and user-specific connections. - Introduced a refreshTokens function to facilitate token refresh logic. - Enhanced MCPTokenStorage to manage client information and refresh token processes. - Updated logging for better traceability during token operations. * chore: cleanup @librechat/auth * feat: implement MCP server initialization in a separate service - Added a new service to handle the initialization of MCP servers, improving code organization and readability. - Refactored the server startup logic to utilize the new initializeMCP function. - Removed redundant MCP initialization code from the main server file. * fix: don't log auth url for user connections * feat: enhance OAuth flow with success and error handling components - Updated OAuth callback routes to redirect to new success and error pages instead of sending status messages. - Introduced `OAuthSuccess` and `OAuthError` components to provide user feedback during authentication. - Added localization support for success and error messages in the translation files. - Implemented countdown functionality in the success component for a better user experience. * fix: refresh token handling for user connections, add missing URL and methods - add standard enum for system user id and helper for determining app-lvel vs. user-level connections * refactor: update token handling in MCPManager and MCPTokenStorage * fix: improve error logging in OAuth authentication handler * fix: concurrency issues for both login url emission and concurrency of oauth flows for shared flows (same user, same server, multiple calls for same server) * fix: properly fail shared flows for concurrent server calls and prevent duplication of tokens * chore: remove unused auth package directory from update configuration * ci: fix mocks in samlStrategy tests * ci: add mcpConfig to AppService test setup * chore: remove obsolete MCP OAuth implementation documentation * fix: update build script for API to use correct command * chore: bump version of @librechat/api to 1.2.4 * fix: update abort signal handling in createMCPTool function * fix: add optional clientInfo parameter to refreshTokensFunction metadata * refactor: replace app.locals.availableTools with getCachedTools in multiple services and controllers for improved tool management * fix: concurrent refresh token handling issue * refactor: add signal parameter to getUserConnection method for improved abort handling * chore: JSDoc typing for `loadEphemeralAgent` * refactor: update isConnectionActive method to use destructured parameters for improved readability * feat: implement caching for MCP tools to handle app-level disconnects for loading list of tools * ci: fix agent test
2025-12-17 00:40:14 +01:00 · 2025-06-17 13:50:33 -04:00 · 2025-06-17 13:50:33 -04:00 · ec7370dfe9
commit ec7370dfe9
parent b412455e9d
60 changed files with 3399 additions and 764 deletions
--- a/packages/api/src/oauth/tokens.ts
+++ b/packages/api/src/oauth/tokens.ts
@ -0,0 +1,324 @@
+import axios from 'axios';
+import { logger } from '@librechat/data-schemas';
+import { TokenExchangeMethodEnum } from 'librechat-data-provider';
+import type { TokenMethods } from '@librechat/data-schemas';
+import type { AxiosError } from 'axios';
+import { encryptV2, decryptV2 } from '~/crypto';
+import { logAxiosError } from '~/utils';
+
+export function createHandleOAuthToken({
+  findToken,
+  updateToken,
+  createToken,
+}: {
+  findToken: TokenMethods['findToken'];
+  updateToken: TokenMethods['updateToken'];
+  createToken: TokenMethods['createToken'];
+}) {
+  /**
+   * Handles the OAuth token by creating or updating the token.
+   * @param fields
+   * @param fields.userId - The user's ID.
+   * @param fields.token - The full token to store.
+   * @param fields.identifier - Unique, alternative identifier for the token.
+   * @param fields.expiresIn - The number of seconds until the token expires.
+   * @param fields.metadata - Additional metadata to store with the token.
+   * @param [fields.type="oauth"] - The type of token. Default is 'oauth'.
+   */
+  return async function handleOAuthToken({
+    token,
+    userId,
+    identifier,
+    expiresIn,
+    metadata,
+    type = 'oauth',
+  }: {
+    token: string;
+    userId: string;
+    identifier: string;
+    expiresIn?: number | string | null;
+    metadata?: Record<string, unknown>;
+    type?: string;
+  }) {
+    const encrypedToken = await encryptV2(token);
+    let expiresInNumber = 3600;
+    if (typeof expiresIn === 'number') {
+      expiresInNumber = expiresIn;
+    } else if (expiresIn != null) {
+      expiresInNumber = parseInt(expiresIn, 10) || 3600;
+    }
+    const tokenData = {
+      type,
+      userId,
+      metadata,
+      identifier,
+      token: encrypedToken,
+      expiresIn: expiresInNumber,
+    };
+
+    const existingToken = await findToken({ userId, identifier });
+    if (existingToken) {
+      return await updateToken({ identifier }, tokenData);
+    } else {
+      return await createToken(tokenData);
+    }
+  };
+}
+
+/**
+ * Processes the access tokens and stores them in the database.
+ * @param tokenData
+ * @param tokenData.access_token
+ * @param tokenData.expires_in
+ * @param [tokenData.refresh_token]
+ * @param [tokenData.refresh_token_expires_in]
+ * @param metadata
+ * @param metadata.userId
+ * @param metadata.identifier
+ */
+async function processAccessTokens(
+  tokenData: {
+    access_token: string;
+    expires_in: number;
+    refresh_token?: string;
+    refresh_token_expires_in?: number;
+  },
+  { userId, identifier }: { userId: string; identifier: string },
+  {
+    findToken,
+    updateToken,
+    createToken,
+  }: {
+    findToken: TokenMethods['findToken'];
+    updateToken: TokenMethods['updateToken'];
+    createToken: TokenMethods['createToken'];
+  },
+) {
+  const { access_token, expires_in = 3600, refresh_token, refresh_token_expires_in } = tokenData;
+  if (!access_token) {
+    logger.error('Access token not found: ', tokenData);
+    throw new Error('Access token not found');
+  }
+  const handleOAuthToken = createHandleOAuthToken({
+    findToken,
+    updateToken,
+    createToken,
+  });
+  await handleOAuthToken({
+    identifier,
+    token: access_token,
+    expiresIn: expires_in,
+    userId,
+  });
+
+  if (refresh_token != null) {
+    logger.debug('Processing refresh token');
+    await handleOAuthToken({
+      token: refresh_token,
+      type: 'oauth_refresh',
+      userId,
+      identifier: `${identifier}:refresh`,
+      expiresIn: refresh_token_expires_in ?? null,
+    });
+  }
+  logger.debug('Access tokens processed');
+}
+
+/**
+ * Refreshes the access token using the refresh token.
+ * @param fields
+ * @param fields.userId - The ID of the user.
+ * @param fields.client_url - The URL of the OAuth provider.
+ * @param fields.identifier - The identifier for the token.
+ * @param fields.refresh_token - The refresh token to use.
+ * @param fields.token_exchange_method - The token exchange method ('default_post' or 'basic_auth_header').
+ * @param fields.encrypted_oauth_client_id - The client ID for the OAuth provider.
+ * @param fields.encrypted_oauth_client_secret - The client secret for the OAuth provider.
+ */
+export async function refreshAccessToken(
+  {
+    userId,
+    client_url,
+    identifier,
+    refresh_token,
+    token_exchange_method,
+    encrypted_oauth_client_id,
+    encrypted_oauth_client_secret,
+  }: {
+    userId: string;
+    client_url: string;
+    identifier: string;
+    refresh_token: string;
+    token_exchange_method: TokenExchangeMethodEnum;
+    encrypted_oauth_client_id: string;
+    encrypted_oauth_client_secret: string;
+  },
+  {
+    findToken,
+    updateToken,
+    createToken,
+  }: {
+    findToken: TokenMethods['findToken'];
+    updateToken: TokenMethods['updateToken'];
+    createToken: TokenMethods['createToken'];
+  },
+): Promise<{
+  access_token: string;
+  expires_in: number;
+  refresh_token?: string;
+  refresh_token_expires_in?: number;
+}> {
+  try {
+    const oauth_client_id = await decryptV2(encrypted_oauth_client_id);
+    const oauth_client_secret = await decryptV2(encrypted_oauth_client_secret);
+
+    const headers: Record<string, string> = {
+      'Content-Type': 'application/x-www-form-urlencoded',
+      Accept: 'application/json',
+    };
+
+    const params = new URLSearchParams({
+      grant_type: 'refresh_token',
+      refresh_token,
+    });
+
+    if (token_exchange_method === TokenExchangeMethodEnum.BasicAuthHeader) {
+      const basicAuth = Buffer.from(`${oauth_client_id}:${oauth_client_secret}`).toString('base64');
+      headers['Authorization'] = `Basic ${basicAuth}`;
+    } else {
+      params.append('client_id', oauth_client_id);
+      params.append('client_secret', oauth_client_secret);
+    }
+
+    const response = await axios({
+      method: 'POST',
+      url: client_url,
+      headers,
+      data: params.toString(),
+    });
+    await processAccessTokens(
+      response.data,
+      {
+        userId,
+        identifier,
+      },
+      {
+        findToken,
+        updateToken,
+        createToken,
+      },
+    );
+    logger.debug(`Access token refreshed successfully for ${identifier}`);
+    return response.data;
+  } catch (error) {
+    const message = 'Error refreshing OAuth tokens';
+    throw new Error(
+      logAxiosError({
+        message,
+        error: error as AxiosError,
+      }),
+    );
+  }
+}
+
+/**
+ * Handles the OAuth callback and exchanges the authorization code for tokens.
+ * @param {object} fields
+ * @param {string} fields.code - The authorization code returned by the provider.
+ * @param {string} fields.userId - The ID of the user.
+ * @param {string} fields.identifier - The identifier for the token.
+ * @param {string} fields.client_url - The URL of the OAuth provider.
+ * @param {string} fields.redirect_uri - The redirect URI for the OAuth provider.
+ * @param {string} fields.token_exchange_method - The token exchange method ('default_post' or 'basic_auth_header').
+ * @param {string} fields.encrypted_oauth_client_id - The client ID for the OAuth provider.
+ * @param {string} fields.encrypted_oauth_client_secret - The client secret for the OAuth provider.
+ */
+export async function getAccessToken(
+  {
+    code,
+    userId,
+    identifier,
+    client_url,
+    redirect_uri,
+    token_exchange_method,
+    encrypted_oauth_client_id,
+    encrypted_oauth_client_secret,
+  }: {
+    code: string;
+    userId: string;
+    identifier: string;
+    client_url: string;
+    redirect_uri: string;
+    token_exchange_method: TokenExchangeMethodEnum;
+    encrypted_oauth_client_id: string;
+    encrypted_oauth_client_secret: string;
+  },
+  {
+    findToken,
+    updateToken,
+    createToken,
+  }: {
+    findToken: TokenMethods['findToken'];
+    updateToken: TokenMethods['updateToken'];
+    createToken: TokenMethods['createToken'];
+  },
+): Promise<{
+  access_token: string;
+  expires_in: number;
+  refresh_token?: string;
+  refresh_token_expires_in?: number;
+}> {
+  const oauth_client_id = await decryptV2(encrypted_oauth_client_id);
+  const oauth_client_secret = await decryptV2(encrypted_oauth_client_secret);
+
+  const headers: Record<string, string> = {
+    'Content-Type': 'application/x-www-form-urlencoded',
+    Accept: 'application/json',
+  };
+
+  const params = new URLSearchParams({
+    code,
+    grant_type: 'authorization_code',
+    redirect_uri,
+  });
+
+  if (token_exchange_method === TokenExchangeMethodEnum.BasicAuthHeader) {
+    const basicAuth = Buffer.from(`${oauth_client_id}:${oauth_client_secret}`).toString('base64');
+    headers['Authorization'] = `Basic ${basicAuth}`;
+  } else {
+    params.append('client_id', oauth_client_id);
+    params.append('client_secret', oauth_client_secret);
+  }
+
+  try {
+    const response = await axios({
+      method: 'POST',
+      url: client_url,
+      headers,
+      data: params.toString(),
+    });
+
+    await processAccessTokens(
+      response.data,
+      {
+        userId,
+        identifier,
+      },
+      {
+        findToken,
+        updateToken,
+        createToken,
+      },
+    );
+    logger.debug(`Access tokens successfully created for ${identifier}`);
+    return response.data;
+  } catch (error) {
+    const message = 'Error exchanging OAuth code';
+    throw new Error(
+      logAxiosError({
+        message,
+        error: error as AxiosError,
+      }),
+    );
+  }
+}