🚦 feat: Auto-reinitialize MCP Servers on Request (#9226)

This commit is contained in:
Danny Avila 2025-08-23 03:27:05 -04:00
parent ac608ded46
commit c827fdd10e
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
28 changed files with 871 additions and 312 deletions

View file

@ -19,7 +19,10 @@ const AgentClient = require('~/server/controllers/agents/client');
const { getAgent } = require('~/models/Agent');
const { logViolation } = require('~/cache');
function createToolLoader() {
/**
* @param {AbortSignal} signal
*/
function createToolLoader(signal) {
/**
* @param {object} params
* @param {ServerRequest} params.req
@ -29,7 +32,11 @@ function createToolLoader() {
* @param {string} params.provider
* @param {string} params.model
* @param {AgentToolResources} params.tool_resources
* @returns {Promise<{ tools: StructuredTool[], toolContextMap: Record<string, unknown> } | undefined>}
* @returns {Promise<{
* tools: StructuredTool[],
* toolContextMap: Record<string, unknown>,
* userMCPAuthMap?: Record<string, Record<string, string>>
* } | undefined>}
*/
return async function loadTools({ req, res, agentId, tools, provider, model, tool_resources }) {
const agent = { id: agentId, tools, provider, model };
@ -38,6 +45,7 @@ function createToolLoader() {
req,
res,
agent,
signal,
tool_resources,
});
} catch (error) {
@ -46,7 +54,7 @@ function createToolLoader() {
};
}
const initializeClient = async ({ req, res, endpointOption }) => {
const initializeClient = async ({ req, res, signal, endpointOption }) => {
if (!endpointOption) {
throw new Error('Endpoint option not provided');
}
@ -92,7 +100,7 @@ const initializeClient = async ({ req, res, endpointOption }) => {
/** @type {Set<string>} */
const allowedProviders = new Set(req?.app?.locals?.[EModelEndpoint.agents]?.allowedProviders);
const loadTools = createToolLoader();
const loadTools = createToolLoader(signal);
/** @type {Array<MongoFile>} */
const requestFiles = req.body.files ?? [];
/** @type {string} */
@ -111,6 +119,7 @@ const initializeClient = async ({ req, res, endpointOption }) => {
});
const agent_ids = primaryConfig.agent_ids;
let userMCPAuthMap = primaryConfig.userMCPAuthMap;
if (agent_ids?.length) {
for (const agentId of agent_ids) {
const agent = await getAgent({ id: agentId });
@ -140,6 +149,7 @@ const initializeClient = async ({ req, res, endpointOption }) => {
endpointOption,
allowedProviders,
});
Object.assign(userMCPAuthMap, config.userMCPAuthMap ?? {});
agentConfigs.set(agentId, config);
}
}
@ -188,7 +198,7 @@ const initializeClient = async ({ req, res, endpointOption }) => {
: EModelEndpoint.agents,
});
return { client };
return { client, userMCPAuthMap };
};
module.exports = { initializeClient };