Andreas/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2025-12-20 18:30:15 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

🛠️ refactor: Consolidate MCP Tool Caching (#9172)

Browse source 
* 🛠️ refactor: Consolidate MCP Tool Caching

* 🐍 fix: Correctly mock and utilize updateMCPUserTools in MCP route tests
This commit is contained in:
Danny Avila 2025-08-20 12:19:29 -04:00 • committed by GitHub
parent 5a14ee9c6a
commit da4aa37493
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 293 additions and 191 deletions
Download patch file Download diff file Expand all files Collapse all files

22
api/server/routes/__tests__/mcp.spec.js
View file

@ -1,9 +1,10 @@
const { MongoMemoryServer } = require('mongodb-memory-server');
const express = require('express');
const request = require('supertest');
const mongoose = require('mongoose');
const express = require('express');
const { MongoMemoryServer } = require('mongodb-memory-server');
jest.mock('@librechat/api', () => ({
...jest.requireActual('@librechat/api'),
MCPOAuthHandler: {
initiateOAuthFlow: jest.fn(),
getFlowState: jest.fn(),
@ -44,6 +45,10 @@ jest.mock('~/server/services/Config', () => ({
loadCustomConfig: jest.fn(),
}));
jest.mock('~/server/services/Config/mcpToolsCache', () => ({
updateMCPUserTools: jest.fn(),
}));
jest.mock('~/server/services/MCP', () => ({
getMCPSetupData: jest.fn(),
getServerConnectionStatus: jest.fn(),
@ -759,8 +764,10 @@ describe('MCP Routes', () => {
require('~/cache').getLogStores.mockReturnValue({});
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
const { updateMCPUserTools } = require('~/server/services/Config/mcpToolsCache');
getCachedTools.mockResolvedValue({});
setCachedTools.mockResolvedValue();
updateMCPUserTools.mockResolvedValue();
const response = await request(app).post('/api/mcp/test-server/reinitialize');
@ -776,7 +783,14 @@ describe('MCP Routes', () => {
'test-user-id',
'test-server',
);
expect(setCachedTools).toHaveBeenCalled();
expect(updateMCPUserTools).toHaveBeenCalledWith({
userId: 'test-user-id',
serverName: 'test-server',
tools: [
{ name: 'tool1', description: 'Test tool 1', inputSchema: { type: 'object' } },
{ name: 'tool2', description: 'Test tool 2', inputSchema: { type: 'object' } },
],
});
});
it('should handle server with custom user variables', async () => {
@ -803,8 +817,10 @@ describe('MCP Routes', () => {
);
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
const { updateMCPUserTools } = require('~/server/services/Config/mcpToolsCache');
getCachedTools.mockResolvedValue({});
setCachedTools.mockResolvedValue();
updateMCPUserTools.mockResolvedValue();
const response = await request(app).post('/api/mcp/test-server/reinitialize');

60
api/server/routes/mcp.js
View file

@ -3,7 +3,7 @@ const { MCPOAuthHandler } = require('@librechat/api');
const { Router } = require('express');
const { getMCPSetupData, getServerConnectionStatus } = require('~/server/services/MCP');
const { findToken, updateToken, createToken, deleteTokens } = require('~/models');
const { setCachedTools, getCachedTools } = require('~/server/services/Config');
const { updateMCPUserTools } = require('~/server/services/Config/mcpToolsCache');
const { getUserPluginAuthValue } = require('~/server/services/PluginService');
const { CacheKeys, Constants } = require('librechat-data-provider');
const { getMCPManager, getFlowStateManager } = require('~/config');
@ -142,33 +142,12 @@ router.get('/:serverName/oauth/callback', async (req, res) => {
`[MCP OAuth] Successfully reconnected ${serverName} for user ${flowState.userId}`,
);
const userTools = (await getCachedTools({ userId: flowState.userId })) || {};
const mcpDelimiter = Constants.mcp_delimiter;
for (const key of Object.keys(userTools)) {
if (key.endsWith(`${mcpDelimiter}${serverName}`)) {
delete userTools[key];
}
}
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,
},
};
}
await setCachedTools(userTools, { userId: flowState.userId });
logger.debug(
`[MCP OAuth] Cached ${tools.length} tools for ${serverName} user ${flowState.userId}`,
);
await updateMCPUserTools({
userId: flowState.userId,
serverName,
tools,
});
} else {
logger.debug(`[MCP OAuth] System-level OAuth completed for ${serverName}`);
}
@ -396,29 +375,12 @@ router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => {
}
if (userConnection && !oauthRequired) {
const userTools = (await getCachedTools({ userId: user.id })) || {};
const mcpDelimiter = Constants.mcp_delimiter;
for (const key of Object.keys(userTools)) {
if (key.endsWith(`${mcpDelimiter}${serverName}`)) {
delete userTools[key];
}
}
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,
},
};
}
await setCachedTools(userTools, { userId: user.id });
await updateMCPUserTools({
userId: user.id,
serverName,
tools,
});
}
logger.debug(