LibreChat/api/server/services/ActionService.js
Danny Avila ec7370dfe9
🪐 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-06-17 13:50:33 -04:00

449 lines
16 KiB
JavaScript

const jwt = require('jsonwebtoken');
const { nanoid } = require('nanoid');
const { tool } = require('@langchain/core/tools');
const { logger } = require('@librechat/data-schemas');
const { GraphEvents, sleep } = require('@librechat/agents');
const {
sendEvent,
encryptV2,
decryptV2,
logAxiosError,
refreshAccessToken,
} = require('@librechat/api');
const {
Time,
CacheKeys,
StepTypes,
Constants,
AuthTypeEnum,
actionDelimiter,
isImageVisionTool,
actionDomainSeparator,
} = require('librechat-data-provider');
const { findToken, updateToken, createToken } = require('~/models');
const { getActions, deleteActions } = require('~/models/Action');
const { deleteAssistant } = require('~/models/Assistant');
const { getFlowStateManager } = require('~/config');
const { getLogStores } = require('~/cache');
const JWT_SECRET = process.env.JWT_SECRET;
const toolNameRegex = /^[a-zA-Z0-9_-]+$/;
const replaceSeparatorRegex = new RegExp(actionDomainSeparator, 'g');
/**
* Validates tool name against regex pattern and updates if necessary.
* @param {object} params - The parameters for the function.
* @param {object} params.req - Express Request.
* @param {FunctionTool} params.tool - The tool object.
* @param {string} params.assistant_id - The assistant ID
* @returns {object|null} - Updated tool object or null if invalid and not an action.
*/
const validateAndUpdateTool = async ({ req, tool, assistant_id }) => {
let actions;
if (isImageVisionTool(tool)) {
return null;
}
if (!toolNameRegex.test(tool.function.name)) {
const [functionName, domain] = tool.function.name.split(actionDelimiter);
actions = await getActions({ assistant_id, user: req.user.id }, true);
const matchingActions = actions.filter((action) => {
const metadata = action.metadata;
return metadata && metadata.domain === domain;
});
const action = matchingActions[0];
if (!action) {
return null;
}
const parsedDomain = await domainParser(domain, true);
if (!parsedDomain) {
return null;
}
tool.function.name = `${functionName}${actionDelimiter}${parsedDomain}`;
}
return tool;
};
/**
* Encodes or decodes a domain name to/from base64, or replacing periods with a custom separator.
*
* Necessary due to `[a-zA-Z0-9_-]*` Regex Validation, limited to a 64-character maximum.
*
* @param {string} domain - The domain name to encode/decode.
* @param {boolean} inverse - False to decode from base64, true to encode to base64.
* @returns {Promise<string>} Encoded or decoded domain string.
*/
async function domainParser(domain, inverse = false) {
if (!domain) {
return;
}
const domainsCache = getLogStores(CacheKeys.ENCODED_DOMAINS);
const cachedDomain = await domainsCache.get(domain);
if (inverse && cachedDomain) {
return domain;
}
if (inverse && domain.length <= Constants.ENCODED_DOMAIN_LENGTH) {
return domain.replace(/\./g, actionDomainSeparator);
}
if (inverse) {
const modifiedDomain = Buffer.from(domain).toString('base64');
const key = modifiedDomain.substring(0, Constants.ENCODED_DOMAIN_LENGTH);
await domainsCache.set(key, modifiedDomain);
return key;
}
if (!cachedDomain) {
return domain.replace(replaceSeparatorRegex, '.');
}
try {
return Buffer.from(cachedDomain, 'base64').toString('utf-8');
} catch (error) {
logger.error(`Failed to parse domain (possibly not base64): ${domain}`, error);
return domain;
}
}
/**
* Loads action sets based on the user and assistant ID.
*
* @param {Object} searchParams - The parameters for loading action sets.
* @param {string} searchParams.user - The user identifier.
* @param {string} [searchParams.agent_id]- The agent identifier.
* @param {string} [searchParams.assistant_id]- The assistant identifier.
* @returns {Promise<Action[] | null>} A promise that resolves to an array of actions or `null` if no match.
*/
async function loadActionSets(searchParams) {
return await getActions(searchParams, true);
}
/**
* Creates a general tool for an entire action set.
*
* @param {Object} params - The parameters for loading action sets.
* @param {string} params.userId
* @param {ServerResponse} params.res
* @param {Action} params.action - The action set. Necessary for decrypting authentication values.
* @param {ActionRequest} params.requestBuilder - The ActionRequest builder class to execute the API call.
* @param {string | undefined} [params.name] - The name of the tool.
* @param {string | undefined} [params.description] - The description for the tool.
* @param {import('zod').ZodTypeAny | undefined} [params.zodSchema] - The Zod schema for tool input validation/definition
* @param {{ oauth_client_id?: string; oauth_client_secret?: string; }} params.encrypted - The encrypted values for the action.
* @returns { Promise<typeof tool | { _call: (toolInput: Object | string) => unknown}> } An object with `_call` method to execute the tool input.
*/
async function createActionTool({
userId,
res,
action,
requestBuilder,
zodSchema,
name,
description,
encrypted,
}) {
/** @type {(toolInput: Object | string, config: GraphRunnableConfig) => Promise<unknown>} */
const _call = async (toolInput, config) => {
try {
/** @type {import('librechat-data-provider').ActionMetadataRuntime} */
const metadata = action.metadata;
const executor = requestBuilder.createExecutor();
const preparedExecutor = executor.setParams(toolInput ?? {});
if (metadata.auth && metadata.auth.type !== AuthTypeEnum.None) {
try {
if (metadata.auth.type === AuthTypeEnum.OAuth && metadata.auth.authorization_url) {
const action_id = action.action_id;
const identifier = `${userId}:${action.action_id}`;
const requestLogin = async () => {
const { args: _args, stepId, ...toolCall } = config.toolCall ?? {};
if (!stepId) {
throw new Error('Tool call is missing stepId');
}
const statePayload = {
nonce: nanoid(),
user: userId,
action_id,
};
const stateToken = jwt.sign(statePayload, JWT_SECRET, { expiresIn: '10m' });
try {
const redirectUri = `${process.env.DOMAIN_CLIENT}/api/actions/${action_id}/oauth/callback`;
const params = new URLSearchParams({
client_id: metadata.oauth_client_id,
scope: metadata.auth.scope,
redirect_uri: redirectUri,
access_type: 'offline',
response_type: 'code',
state: stateToken,
});
const authURL = `${metadata.auth.authorization_url}?${params.toString()}`;
/** @type {{ id: string; delta: AgentToolCallDelta }} */
const data = {
id: stepId,
delta: {
type: StepTypes.TOOL_CALLS,
tool_calls: [{ ...toolCall, args: '' }],
auth: authURL,
expires_at: Date.now() + Time.TWO_MINUTES,
},
};
const flowsCache = getLogStores(CacheKeys.FLOWS);
const flowManager = getFlowStateManager(flowsCache);
await flowManager.createFlowWithHandler(
`${identifier}:oauth_login:${config.metadata.thread_id}:${config.metadata.run_id}`,
'oauth_login',
async () => {
sendEvent(res, { event: GraphEvents.ON_RUN_STEP_DELTA, data });
logger.debug('Sent OAuth login request to client', { action_id, identifier });
return true;
},
config?.signal,
);
logger.debug('Waiting for OAuth Authorization response', { action_id, identifier });
const result = await flowManager.createFlow(
identifier,
'oauth',
{
state: stateToken,
userId: userId,
client_url: metadata.auth.client_url,
redirect_uri: `${process.env.DOMAIN_SERVER}/api/actions/${action_id}/oauth/callback`,
token_exchange_method: metadata.auth.token_exchange_method,
/** Encrypted values */
encrypted_oauth_client_id: encrypted.oauth_client_id,
encrypted_oauth_client_secret: encrypted.oauth_client_secret,
},
config?.signal,
);
logger.debug('Received OAuth Authorization response', { action_id, identifier });
data.delta.auth = undefined;
data.delta.expires_at = undefined;
sendEvent(res, { event: GraphEvents.ON_RUN_STEP_DELTA, data });
await sleep(3000);
metadata.oauth_access_token = result.access_token;
metadata.oauth_refresh_token = result.refresh_token;
const expiresAt = new Date(Date.now() + result.expires_in * 1000);
metadata.oauth_token_expires_at = expiresAt.toISOString();
} catch (error) {
const errorMessage = 'Failed to authenticate OAuth tool';
logger.error(errorMessage, error);
throw new Error(errorMessage);
}
};
const tokenPromises = [];
tokenPromises.push(findToken({ userId, type: 'oauth', identifier }));
tokenPromises.push(
findToken({
userId,
type: 'oauth_refresh',
identifier: `${identifier}:refresh`,
}),
);
const [tokenData, refreshTokenData] = await Promise.all(tokenPromises);
if (tokenData) {
// Valid token exists, add it to metadata for setAuth
metadata.oauth_access_token = await decryptV2(tokenData.token);
if (refreshTokenData) {
metadata.oauth_refresh_token = await decryptV2(refreshTokenData.token);
}
metadata.oauth_token_expires_at = tokenData.expiresAt.toISOString();
} else if (!refreshTokenData) {
// No tokens exist, need to authenticate
await requestLogin();
} else if (refreshTokenData) {
// Refresh token is still valid, use it to get new access token
try {
const refresh_token = await decryptV2(refreshTokenData.token);
const refreshTokens = async () =>
await refreshAccessToken(
{
userId,
identifier,
refresh_token,
client_url: metadata.auth.client_url,
encrypted_oauth_client_id: encrypted.oauth_client_id,
token_exchange_method: metadata.auth.token_exchange_method,
encrypted_oauth_client_secret: encrypted.oauth_client_secret,
},
{
findToken,
updateToken,
createToken,
},
);
const flowsCache = getLogStores(CacheKeys.FLOWS);
const flowManager = getFlowStateManager(flowsCache);
const refreshData = await flowManager.createFlowWithHandler(
`${identifier}:refresh`,
'oauth_refresh',
refreshTokens,
config?.signal,
);
metadata.oauth_access_token = refreshData.access_token;
if (refreshData.refresh_token) {
metadata.oauth_refresh_token = refreshData.refresh_token;
}
const expiresAt = new Date(Date.now() + refreshData.expires_in * 1000);
metadata.oauth_token_expires_at = expiresAt.toISOString();
} catch (error) {
logger.error('Failed to refresh token, requesting new login:', error);
await requestLogin();
}
} else {
await requestLogin();
}
}
await preparedExecutor.setAuth(metadata);
} catch (error) {
if (
error.message.includes('No access token found') ||
error.message.includes('Access token is expired')
) {
throw error;
}
throw new Error(`Authentication failed: ${error.message}`);
}
}
const response = await preparedExecutor.execute();
if (typeof response.data === 'object') {
return JSON.stringify(response.data);
}
return response.data;
} catch (error) {
const message = `API call to ${action.metadata.domain} failed:`;
return logAxiosError({ message, error });
}
};
if (name) {
return tool(_call, {
name: name.replace(replaceSeparatorRegex, '_'),
description: description || '',
schema: zodSchema,
});
}
return {
_call,
};
}
/**
* Encrypts a sensitive value.
* @param {string} value
* @returns {Promise<string>}
*/
async function encryptSensitiveValue(value) {
// Encode API key to handle special characters like ":"
const encodedValue = encodeURIComponent(value);
return await encryptV2(encodedValue);
}
/**
* Decrypts a sensitive value.
* @param {string} value
* @returns {Promise<string>}
*/
async function decryptSensitiveValue(value) {
const decryptedValue = await decryptV2(value);
return decodeURIComponent(decryptedValue);
}
/**
* Encrypts sensitive metadata values for an action.
*
* @param {ActionMetadata} metadata - The action metadata to encrypt.
* @returns {Promise<ActionMetadata>} The updated action metadata with encrypted values.
*/
async function encryptMetadata(metadata) {
const encryptedMetadata = { ...metadata };
// ServiceHttp
if (metadata.auth && metadata.auth.type === AuthTypeEnum.ServiceHttp) {
if (metadata.api_key) {
encryptedMetadata.api_key = await encryptSensitiveValue(metadata.api_key);
}
}
// OAuth
else if (metadata.auth && metadata.auth.type === AuthTypeEnum.OAuth) {
if (metadata.oauth_client_id) {
encryptedMetadata.oauth_client_id = await encryptSensitiveValue(metadata.oauth_client_id);
}
if (metadata.oauth_client_secret) {
encryptedMetadata.oauth_client_secret = await encryptSensitiveValue(
metadata.oauth_client_secret,
);
}
}
return encryptedMetadata;
}
/**
* Decrypts sensitive metadata values for an action.
*
* @param {ActionMetadata} metadata - The action metadata to decrypt.
* @returns {Promise<ActionMetadata>} The updated action metadata with decrypted values.
*/
async function decryptMetadata(metadata) {
const decryptedMetadata = { ...metadata };
// ServiceHttp
if (metadata.auth && metadata.auth.type === AuthTypeEnum.ServiceHttp) {
if (metadata.api_key) {
decryptedMetadata.api_key = await decryptSensitiveValue(metadata.api_key);
}
}
// OAuth
else if (metadata.auth && metadata.auth.type === AuthTypeEnum.OAuth) {
if (metadata.oauth_client_id) {
decryptedMetadata.oauth_client_id = await decryptSensitiveValue(metadata.oauth_client_id);
}
if (metadata.oauth_client_secret) {
decryptedMetadata.oauth_client_secret = await decryptSensitiveValue(
metadata.oauth_client_secret,
);
}
}
return decryptedMetadata;
}
/**
* Deletes an action and its corresponding assistant.
* @param {Object} params - The parameters for the function.
* @param {OpenAIClient} params.req - The Express Request object.
* @param {string} params.assistant_id - The ID of the assistant.
*/
const deleteAssistantActions = async ({ req, assistant_id }) => {
try {
await deleteActions({ assistant_id, user: req.user.id });
await deleteAssistant({ assistant_id, user: req.user.id });
} catch (error) {
const message = 'Trouble deleting Assistant Actions for Assistant ID: ' + assistant_id;
logger.error(message, error);
throw new Error(message);
}
};
module.exports = {
deleteAssistantActions,
validateAndUpdateTool,
createActionTool,
encryptMetadata,
decryptMetadata,
loadActionSets,
domainParser,
};