mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-28 04:37:37 +00:00
📡 fix: Refresh MCP Tools After List-Changed Notifications (#14686)
* fix(mcp): handle dynamic tool list changes Co-authored-by: Pascal Garber <pascal@artandcode.studio> * test(mcp): fix CI validation * fix(mcp): keep dynamic tool catalogs live * fix(mcp): harden dynamic catalog lifecycle * test(mcp): use typed startup connection * test(mcp): isolate dynamic e2e fixtures * fix(mcp): refresh tools after reconnect * fix(mcp): close dynamic catalog cache gaps * test(mcp): update OAuth connection mocks * fix(mcp): preserve app snapshot ownership * style(mcp): sort connection imports * fix(mcp): close review race conditions * fix(mcp): preserve cache ownership edges * fix(mcp): harden recovery lifecycle * fix(mcp): guard tool-less app refresh * fix(mcp): fence distributed cache races * fix(mcp): retire stale connection state * fix(mcp): keep tool snapshots authoritative * fix(mcp): fence stale app tool publications * style(mcp): sort repository test imports * test(mcp): mock empty startup publication * fix(mcp): preserve app publication generations * fix(mcp): harden publication recovery races * fix(mcp): address tool catalogs by runtime config * fix(mcp): load scoped catalogs for assistant writes * fix(mcp): harden catalog publication recovery * fix(mcp): serialize forced connection replacement * fix(mcp): serialize ordinary creation with replacements * fix(mcp): harden catalog fallback boundaries * fix(mcp): close lifecycle fencing gaps * fix(mcp): preserve catalog authority on failures * fix(mcp): compensate failed catalog mutations * fix(mcp): fence catalog refresh ordering * style(mcp): sort agent loader imports * fix(mcp): cancel stale connection creation * fix(mcp): fence catalog coordination * fix(mcp): close catalog race windows * fix(mcp): harden cross-pod catalog fencing * fix(mcp): close catalog lifecycle edges * style(mcp): sort assistant imports * fix(mcp): reject stale recovery authority * fix(mcp): restore static catalog on every startup * fix(mcp): order app catalog publications * style(mcp): sort catalog revision imports * fix(mcp): separate catalog allocation and commit fences --------- Co-authored-by: Pascal Garber <pascal@artandcode.studio>
This commit is contained in:
Danny Avila
GitHub
parent
ef38f362ec
commit
1bccc2bc18
60 changed files with 8046 additions and 831 deletions
|
|
@ -177,14 +177,24 @@ const deleteUserMcpServers = async (userId) => {
|
|||
const allServersToDelete = [...aclOwnedServers, ...legacyServers];
|
||||
|
||||
const mcpManager = getMCPManager();
|
||||
if (mcpManager) {
|
||||
await Promise.all(
|
||||
allServersToDelete.map(async (s) => {
|
||||
await mcpManager.disconnectUserConnection(userId, s.serverName);
|
||||
await Promise.allSettled(
|
||||
allServersToDelete.map(async (s) => {
|
||||
try {
|
||||
await invalidateCachedTools({ userId, serverName: s.serverName });
|
||||
}),
|
||||
);
|
||||
}
|
||||
} catch (error) {
|
||||
logger.warn(
|
||||
`[deleteUserMcpServers] Failed to invalidate tools for ${s.serverName}:`,
|
||||
error,
|
||||
);
|
||||
} finally {
|
||||
try {
|
||||
await mcpManager?.disconnectUserConnection(userId, s.serverName);
|
||||
} catch (error) {
|
||||
logger.warn(`[deleteUserMcpServers] Failed to disconnect ${s.serverName}:`, error);
|
||||
}
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
await AclEntry.deleteMany({
|
||||
resourceType: ResourceType.MCPSERVER,
|
||||
|
|
@ -295,21 +305,37 @@ const updateUserPluginsController = async (req, res) => {
|
|||
if (pluginKey.startsWith(Constants.mcp_prefix)) {
|
||||
try {
|
||||
const mcpManager = getMCPManager();
|
||||
// Extract server name from pluginKey (format: "mcp_<serverName>")
|
||||
const serverName = pluginKey.replace(Constants.mcp_prefix, '');
|
||||
if (mcpManager) {
|
||||
// Extract server name from pluginKey (format: "mcp_<serverName>")
|
||||
const serverName = pluginKey.replace(Constants.mcp_prefix, '');
|
||||
logger.info(
|
||||
`[updateUserPluginsController] Attempting disconnect of MCP server "${serverName}" for user ${user.id} after plugin auth update.`,
|
||||
);
|
||||
await mcpManager.disconnectUserConnection(user.id, serverName);
|
||||
}
|
||||
let invalidationError;
|
||||
try {
|
||||
await invalidateCachedTools({ userId: user.id, serverName });
|
||||
} catch (error) {
|
||||
invalidationError = error;
|
||||
}
|
||||
try {
|
||||
await mcpManager?.disconnectUserConnection(user.id, serverName);
|
||||
} catch (error) {
|
||||
logger.error(
|
||||
`[updateUserPluginsController] Error disconnecting MCP connection for user ${user.id} after plugin auth update:`,
|
||||
error,
|
||||
);
|
||||
}
|
||||
if (invalidationError) {
|
||||
throw invalidationError;
|
||||
}
|
||||
} catch (disconnectError) {
|
||||
logger.error(
|
||||
`[updateUserPluginsController] Error disconnecting MCP connection for user ${user.id} after plugin auth update:`,
|
||||
`[updateUserPluginsController] Error fencing MCP connection for user ${user.id} after plugin auth update:`,
|
||||
disconnectError,
|
||||
);
|
||||
// Do not fail the request for this, but log it.
|
||||
// A credential mutation is not safely published until the shared generation fence moves.
|
||||
throw disconnectError;
|
||||
}
|
||||
}
|
||||
return res.status(status).send();
|
||||
|
|
|
|||
|
|
@ -174,6 +174,38 @@ beforeEach(() => {
|
|||
});
|
||||
|
||||
describe('updateUserPluginsController MCP OAuth cleanup', () => {
|
||||
it('invalidates the shared tool generation even when local disconnect fails', async () => {
|
||||
const { mcpManager } = setupMCPMocks();
|
||||
mcpManager.disconnectUserConnection.mockRejectedValue(new Error('local dispose failed'));
|
||||
MCPTokenStorage.getClientInfoAndMetadata.mockResolvedValue(null);
|
||||
|
||||
const res = createResponse();
|
||||
await updateUserPluginsController(createRequest(), res);
|
||||
|
||||
expect(mockInvalidateCachedTools).toHaveBeenCalledWith({
|
||||
userId: 'user-1',
|
||||
serverName: 'test-server',
|
||||
});
|
||||
expect(mockInvalidateCachedTools.mock.invocationCallOrder[0]).toBeLessThan(
|
||||
mcpManager.disconnectUserConnection.mock.invocationCallOrder[0],
|
||||
);
|
||||
expect(res.status).toHaveBeenCalledWith(200);
|
||||
});
|
||||
|
||||
it('fails the credential update response when the shared generation fence cannot move', async () => {
|
||||
const { mcpManager } = setupMCPMocks();
|
||||
const fenceError = new Error('Redis unavailable');
|
||||
mockInvalidateCachedTools.mockRejectedValue(fenceError);
|
||||
MCPTokenStorage.getClientInfoAndMetadata.mockResolvedValue(null);
|
||||
|
||||
const res = createResponse();
|
||||
await updateUserPluginsController(createRequest(), res);
|
||||
|
||||
expect(mcpManager.disconnectUserConnection).toHaveBeenCalledWith('user-1', 'test-server');
|
||||
expect(res.status).toHaveBeenCalledWith(500);
|
||||
expect(logger.error).toHaveBeenCalledWith('[updateUserPluginsController]', fenceError);
|
||||
});
|
||||
|
||||
it('clears stored OAuth token state when client metadata is missing', async () => {
|
||||
const { flowManager, mcpManager } = setupMCPMocks();
|
||||
MCPTokenStorage.getClientInfoAndMetadata.mockResolvedValue(null);
|
||||
|
|
|
|||
|
|
@ -130,6 +130,42 @@ describe('deleteUserMcpServers', () => {
|
|||
});
|
||||
});
|
||||
|
||||
test('should delete owned servers when cache invalidation fails', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const server = await MCPServer.create({
|
||||
serverName: 'cache-failure-server',
|
||||
config: { title: 'Cache Failure Server' },
|
||||
author: userId,
|
||||
});
|
||||
|
||||
await permissionService.grantPermission({
|
||||
principalType: PrincipalType.USER,
|
||||
principalId: userId,
|
||||
resourceType: ResourceType.MCPSERVER,
|
||||
resourceId: server._id,
|
||||
accessRoleId: AccessRoleIds.MCPSERVER_OWNER,
|
||||
grantedBy: userId,
|
||||
});
|
||||
|
||||
const disconnectUserConnection = jest.fn().mockResolvedValue(undefined);
|
||||
mockGetMCPManager.mockReturnValue({ disconnectUserConnection });
|
||||
mockInvalidateCachedTools.mockRejectedValueOnce(new Error('Redis unavailable'));
|
||||
|
||||
await deleteUserMcpServers(userId.toString());
|
||||
|
||||
expect(disconnectUserConnection).toHaveBeenCalledWith(
|
||||
userId.toString(),
|
||||
'cache-failure-server',
|
||||
);
|
||||
expect(await MCPServer.findById(server._id)).toBeNull();
|
||||
await expect(
|
||||
AclEntry.countDocuments({
|
||||
resourceType: ResourceType.MCPSERVER,
|
||||
resourceId: server._id,
|
||||
}),
|
||||
).resolves.toBe(0);
|
||||
});
|
||||
|
||||
test('should preserve multi-owned MCP servers', async () => {
|
||||
const deletingUserId = new mongoose.Types.ObjectId();
|
||||
const otherOwnerId = new mongoose.Types.ObjectId();
|
||||
|
|
@ -263,6 +299,10 @@ describe('deleteUserMcpServers', () => {
|
|||
await deleteUserMcpServers(userId.toString());
|
||||
|
||||
expect(await MCPServer.findById(server._id)).toBeNull();
|
||||
expect(mockInvalidateCachedTools).toHaveBeenCalledWith({
|
||||
userId: userId.toString(),
|
||||
serverName: 'no-manager-server',
|
||||
});
|
||||
});
|
||||
|
||||
test('should delete legacy MCP servers that have author but no ACL entries', async () => {
|
||||
|
|
|
|||
|
|
@ -24,11 +24,16 @@ jest.mock('~/server/services/GraphApiService', () => ({
|
|||
|
||||
const mockRegistryInstance = {
|
||||
getServerConfig: jest.fn(),
|
||||
inspectServerUpdate: jest.fn(),
|
||||
commitServerUpdate: jest.fn(),
|
||||
updateServer: jest.fn(),
|
||||
removeServer: jest.fn(),
|
||||
};
|
||||
const mockMcpManager = { disconnectUserConnection: jest.fn() };
|
||||
|
||||
jest.mock('~/config', () => ({
|
||||
logger: { debug: jest.fn(), info: jest.fn(), warn: jest.fn(), error: jest.fn() },
|
||||
getMCPManager: jest.fn(),
|
||||
getMCPManager: jest.fn(() => mockMcpManager),
|
||||
getMCPServersRegistry: jest.fn(() => mockRegistryInstance),
|
||||
}));
|
||||
|
||||
|
|
@ -41,10 +46,17 @@ jest.mock('~/server/services/MCP', () => ({
|
|||
|
||||
jest.mock('~/server/services/Config', () => ({
|
||||
cacheMCPServerTools: jest.fn(),
|
||||
getMCPToolsCacheGeneration: jest.fn().mockResolvedValue('test-generation'),
|
||||
getMCPServerTools: jest.fn(),
|
||||
invalidateCachedTools: jest.fn(),
|
||||
}));
|
||||
|
||||
const { getMCPServersList, getMCPServerById } = require('~/server/controllers/mcp');
|
||||
const {
|
||||
getMCPServersList,
|
||||
getMCPServerById,
|
||||
updateMCPServerController,
|
||||
deleteMCPServerController,
|
||||
} = require('~/server/controllers/mcp');
|
||||
const { grantPermission } = require('~/server/services/PermissionService');
|
||||
const { seedDefaultRoles } = require('~/models');
|
||||
|
||||
|
|
@ -108,6 +120,16 @@ beforeEach(async () => {
|
|||
await User.deleteMany({});
|
||||
mockResolveAllMcpConfigs.mockReset();
|
||||
mockRegistryInstance.getServerConfig.mockReset();
|
||||
mockRegistryInstance.inspectServerUpdate.mockReset();
|
||||
mockRegistryInstance.commitServerUpdate.mockReset();
|
||||
mockRegistryInstance.updateServer.mockReset();
|
||||
mockRegistryInstance.removeServer.mockReset();
|
||||
mockMcpManager.disconnectUserConnection.mockReset().mockResolvedValue(undefined);
|
||||
const cacheService = require('~/server/services/Config');
|
||||
cacheService.invalidateCachedTools.mockReset().mockResolvedValue(undefined);
|
||||
cacheService.getMCPServerTools.mockReset().mockResolvedValue({ retained: {} });
|
||||
cacheService.getMCPToolsCacheGeneration.mockReset().mockResolvedValue('restored-generation');
|
||||
cacheService.cacheMCPServerTools.mockReset().mockResolvedValue(undefined);
|
||||
existsSpy = jest.spyOn(SystemGrant, 'exists');
|
||||
});
|
||||
|
||||
|
|
@ -256,3 +278,212 @@ describe('getMCPServerById', () => {
|
|||
expect(payload.oauth.authorization_url).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe('DB-backed server mutation fencing', () => {
|
||||
const updatedConfig = {
|
||||
type: 'streamable-http',
|
||||
url: 'https://updated.example.com/mcp',
|
||||
source: 'user',
|
||||
};
|
||||
|
||||
it('inspects, fences, commits, fences cross-replica creations, and disconnects', async () => {
|
||||
const user = await createUser();
|
||||
mockRegistryInstance.getServerConfig.mockResolvedValue(
|
||||
createDbConfig(new mongoose.Types.ObjectId()),
|
||||
);
|
||||
mockRegistryInstance.inspectServerUpdate.mockResolvedValue(updatedConfig);
|
||||
mockRegistryInstance.commitServerUpdate.mockResolvedValue(updatedConfig);
|
||||
const res = createRes();
|
||||
|
||||
await updateMCPServerController(
|
||||
{ user, params: { serverName: 'github' }, body: { config: updatedConfig } },
|
||||
res,
|
||||
);
|
||||
|
||||
const { invalidateCachedTools } = require('~/server/services/Config');
|
||||
expect(invalidateCachedTools).toHaveBeenCalledWith({ userId: user.id, serverName: 'github' });
|
||||
expect(invalidateCachedTools).toHaveBeenCalledTimes(2);
|
||||
expect(mockMcpManager.disconnectUserConnection).toHaveBeenCalledWith(user.id, 'github');
|
||||
expect(mockRegistryInstance.inspectServerUpdate.mock.invocationCallOrder[0]).toBeLessThan(
|
||||
invalidateCachedTools.mock.invocationCallOrder[0],
|
||||
);
|
||||
expect(invalidateCachedTools.mock.invocationCallOrder[0]).toBeLessThan(
|
||||
mockRegistryInstance.commitServerUpdate.mock.invocationCallOrder[0],
|
||||
);
|
||||
expect(mockRegistryInstance.commitServerUpdate.mock.invocationCallOrder[0]).toBeLessThan(
|
||||
invalidateCachedTools.mock.invocationCallOrder[1],
|
||||
);
|
||||
expect(invalidateCachedTools.mock.invocationCallOrder[1]).toBeLessThan(
|
||||
mockMcpManager.disconnectUserConnection.mock.invocationCallOrder[0],
|
||||
);
|
||||
expect(res.status).toHaveBeenCalledWith(200);
|
||||
});
|
||||
|
||||
it('does not fence the valid catalog when update inspection or persistence fails', async () => {
|
||||
const user = await createUser();
|
||||
const updateError = new Error('inspection failed');
|
||||
mockRegistryInstance.getServerConfig.mockResolvedValue(
|
||||
createDbConfig(new mongoose.Types.ObjectId()),
|
||||
);
|
||||
mockRegistryInstance.inspectServerUpdate.mockRejectedValue(updateError);
|
||||
const res = createRes();
|
||||
|
||||
await updateMCPServerController(
|
||||
{ user, params: { serverName: 'github' }, body: { config: updatedConfig } },
|
||||
res,
|
||||
);
|
||||
|
||||
expect(require('~/server/services/Config').invalidateCachedTools).not.toHaveBeenCalled();
|
||||
expect(mockRegistryInstance.commitServerUpdate).not.toHaveBeenCalled();
|
||||
expect(mockMcpManager.disconnectUserConnection).not.toHaveBeenCalled();
|
||||
expect(res.status).toHaveBeenCalledWith(500);
|
||||
});
|
||||
|
||||
it('does not commit an inspected update when the distributed fence fails', async () => {
|
||||
const user = await createUser();
|
||||
mockRegistryInstance.getServerConfig.mockResolvedValue(
|
||||
createDbConfig(new mongoose.Types.ObjectId()),
|
||||
);
|
||||
mockRegistryInstance.inspectServerUpdate.mockResolvedValue(updatedConfig);
|
||||
require('~/server/services/Config').invalidateCachedTools.mockRejectedValue(
|
||||
new Error('Redis unavailable'),
|
||||
);
|
||||
const res = createRes();
|
||||
|
||||
await updateMCPServerController(
|
||||
{ user, params: { serverName: 'github' }, body: { config: updatedConfig } },
|
||||
res,
|
||||
);
|
||||
|
||||
expect(mockMcpManager.disconnectUserConnection).not.toHaveBeenCalled();
|
||||
expect(mockRegistryInstance.commitServerUpdate).not.toHaveBeenCalled();
|
||||
expect(res.status).toHaveBeenCalledWith(500);
|
||||
});
|
||||
|
||||
it('restores the retained catalog when update persistence fails after fencing', async () => {
|
||||
const user = await createUser();
|
||||
const existingConfig = createDbConfig(new mongoose.Types.ObjectId());
|
||||
const retainedTools = { retained: { function: { name: 'retained' } } };
|
||||
mockRegistryInstance.getServerConfig.mockResolvedValue(existingConfig);
|
||||
mockRegistryInstance.inspectServerUpdate.mockResolvedValue(updatedConfig);
|
||||
mockRegistryInstance.commitServerUpdate.mockRejectedValue(new Error('database unavailable'));
|
||||
require('~/server/services/Config').getMCPServerTools.mockResolvedValue(retainedTools);
|
||||
const res = createRes();
|
||||
|
||||
await updateMCPServerController(
|
||||
{ user, params: { serverName: 'github' }, body: { config: updatedConfig } },
|
||||
res,
|
||||
);
|
||||
|
||||
expect(require('~/server/services/Config').cacheMCPServerTools).toHaveBeenCalledWith({
|
||||
userId: user.id,
|
||||
serverName: 'github',
|
||||
serverConfig: existingConfig,
|
||||
serverTools: retainedTools,
|
||||
publicationGeneration: 'restored-generation',
|
||||
});
|
||||
expect(mockMcpManager.disconnectUserConnection).not.toHaveBeenCalled();
|
||||
expect(res.status).toHaveBeenCalledWith(500);
|
||||
});
|
||||
|
||||
it('continues an update when only local disconnect cleanup fails', async () => {
|
||||
const user = await createUser();
|
||||
mockRegistryInstance.getServerConfig.mockResolvedValue(
|
||||
createDbConfig(new mongoose.Types.ObjectId()),
|
||||
);
|
||||
mockMcpManager.disconnectUserConnection.mockRejectedValue(new Error('dispose failed'));
|
||||
mockRegistryInstance.inspectServerUpdate.mockResolvedValue(updatedConfig);
|
||||
mockRegistryInstance.commitServerUpdate.mockResolvedValue(updatedConfig);
|
||||
const res = createRes();
|
||||
|
||||
await updateMCPServerController(
|
||||
{ user, params: { serverName: 'github' }, body: { config: updatedConfig } },
|
||||
res,
|
||||
);
|
||||
|
||||
expect(mockRegistryInstance.commitServerUpdate).toHaveBeenCalled();
|
||||
expect(res.status).toHaveBeenCalledWith(200);
|
||||
});
|
||||
|
||||
it('retries a transient post-commit fence failure before returning success', async () => {
|
||||
const user = await createUser();
|
||||
mockRegistryInstance.getServerConfig.mockResolvedValue(
|
||||
createDbConfig(new mongoose.Types.ObjectId()),
|
||||
);
|
||||
mockRegistryInstance.inspectServerUpdate.mockResolvedValue(updatedConfig);
|
||||
mockRegistryInstance.commitServerUpdate.mockResolvedValue(updatedConfig);
|
||||
require('~/server/services/Config')
|
||||
.invalidateCachedTools.mockResolvedValueOnce(undefined)
|
||||
.mockRejectedValueOnce(new Error('Redis MOVED'))
|
||||
.mockResolvedValueOnce(undefined);
|
||||
const res = createRes();
|
||||
|
||||
await updateMCPServerController(
|
||||
{ user, params: { serverName: 'github' }, body: { config: updatedConfig } },
|
||||
res,
|
||||
);
|
||||
|
||||
expect(require('~/server/services/Config').invalidateCachedTools).toHaveBeenCalledTimes(3);
|
||||
expect(mockMcpManager.disconnectUserConnection).toHaveBeenCalledWith(user.id, 'github');
|
||||
expect(res.status).toHaveBeenCalledWith(200);
|
||||
});
|
||||
|
||||
it('fences before deletion and fences cross-replica creations before disconnecting', async () => {
|
||||
const user = await createUser();
|
||||
mockRegistryInstance.removeServer.mockResolvedValue(undefined);
|
||||
const res = createRes();
|
||||
|
||||
await deleteMCPServerController({ user, params: { serverName: 'github' } }, res);
|
||||
|
||||
const { invalidateCachedTools } = require('~/server/services/Config');
|
||||
expect(invalidateCachedTools).toHaveBeenCalledWith({ userId: user.id, serverName: 'github' });
|
||||
expect(invalidateCachedTools).toHaveBeenCalledTimes(2);
|
||||
expect(mockMcpManager.disconnectUserConnection).toHaveBeenCalledWith(user.id, 'github');
|
||||
expect(invalidateCachedTools.mock.invocationCallOrder[0]).toBeLessThan(
|
||||
mockRegistryInstance.removeServer.mock.invocationCallOrder[0],
|
||||
);
|
||||
expect(mockRegistryInstance.removeServer.mock.invocationCallOrder[0]).toBeLessThan(
|
||||
invalidateCachedTools.mock.invocationCallOrder[1],
|
||||
);
|
||||
expect(invalidateCachedTools.mock.invocationCallOrder[1]).toBeLessThan(
|
||||
mockMcpManager.disconnectUserConnection.mock.invocationCallOrder[0],
|
||||
);
|
||||
expect(res.status).toHaveBeenCalledWith(200);
|
||||
});
|
||||
|
||||
it('does not delete the registry entry when the distributed fence fails', async () => {
|
||||
const user = await createUser();
|
||||
require('~/server/services/Config').invalidateCachedTools.mockRejectedValue(
|
||||
new Error('Redis unavailable'),
|
||||
);
|
||||
const res = createRes();
|
||||
|
||||
await deleteMCPServerController({ user, params: { serverName: 'github' } }, res);
|
||||
|
||||
expect(mockMcpManager.disconnectUserConnection).not.toHaveBeenCalled();
|
||||
expect(mockRegistryInstance.removeServer).not.toHaveBeenCalled();
|
||||
expect(res.status).toHaveBeenCalledWith(500);
|
||||
});
|
||||
|
||||
it('restores the retained catalog when deletion persistence fails after fencing', async () => {
|
||||
const user = await createUser();
|
||||
const existingConfig = createDbConfig(new mongoose.Types.ObjectId());
|
||||
const retainedTools = { retained: { function: { name: 'retained' } } };
|
||||
mockRegistryInstance.getServerConfig.mockResolvedValue(existingConfig);
|
||||
mockRegistryInstance.removeServer.mockRejectedValue(new Error('Deletion failed'));
|
||||
require('~/server/services/Config').getMCPServerTools.mockResolvedValue(retainedTools);
|
||||
const res = createRes();
|
||||
|
||||
await deleteMCPServerController({ user, params: { serverName: 'github' } }, res);
|
||||
|
||||
expect(require('~/server/services/Config').cacheMCPServerTools).toHaveBeenCalledWith({
|
||||
userId: user.id,
|
||||
serverName: 'github',
|
||||
serverConfig: existingConfig,
|
||||
serverTools: retainedTools,
|
||||
publicationGeneration: 'restored-generation',
|
||||
});
|
||||
expect(mockMcpManager.disconnectUserConnection).not.toHaveBeenCalled();
|
||||
expect(res.status).toHaveBeenCalledWith(500);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -7,8 +7,7 @@ const validateAuthor = require('~/server/middleware/assistants/validateAuthor');
|
|||
const { getStrategyFunctions } = require('~/server/services/Files/strategies');
|
||||
const { deleteAssistantActions } = require('~/server/services/ActionService');
|
||||
const { getOpenAIClient, fetchAssistants } = require('./helpers');
|
||||
const { healMcpToolNames } = require('~/server/services/MCP');
|
||||
const { getCachedTools } = require('~/server/services/Config');
|
||||
const { healMcpToolNames, getAssistantToolDefinitions } = require('~/server/services/MCP');
|
||||
const { manifestToolMap, isAgentsOnlyTool } = require('~/app/clients/tools');
|
||||
|
||||
/**
|
||||
|
|
@ -31,7 +30,7 @@ const createAssistant = async (req, res) => {
|
|||
delete assistantData.conversation_starters;
|
||||
delete assistantData.append_current_datetime;
|
||||
|
||||
const toolDefinitions = (await getCachedTools()) ?? {};
|
||||
const toolDefinitions = await getAssistantToolDefinitions({ req, tools });
|
||||
const healedTools = await healMcpToolNames({ req, tools, toolDefinitions });
|
||||
|
||||
assistantData.tools = healedTools
|
||||
|
|
@ -146,7 +145,7 @@ const patchAssistant = async (req, res) => {
|
|||
...updateData
|
||||
} = req.body;
|
||||
|
||||
const toolDefinitions = (await getCachedTools()) ?? {};
|
||||
const toolDefinitions = await getAssistantToolDefinitions({ req, tools: updateData.tools });
|
||||
const healedTools = await healMcpToolNames({ req, tools: updateData.tools, toolDefinitions });
|
||||
|
||||
updateData.tools = healedTools
|
||||
|
|
|
|||
|
|
@ -2,8 +2,7 @@ const { logger } = require('@librechat/data-schemas');
|
|||
const { ToolCallTypes } = require('librechat-data-provider');
|
||||
const validateAuthor = require('~/server/middleware/assistants/validateAuthor');
|
||||
const { validateAndUpdateTool } = require('~/server/services/ActionService');
|
||||
const { healMcpToolNames } = require('~/server/services/MCP');
|
||||
const { getCachedTools } = require('~/server/services/Config');
|
||||
const { healMcpToolNames, getAssistantToolDefinitions } = require('~/server/services/MCP');
|
||||
const { manifestToolMap, isAgentsOnlyTool } = require('~/app/clients/tools');
|
||||
const { updateAssistantDoc } = require('~/models');
|
||||
const { getOpenAIClient } = require('./helpers');
|
||||
|
|
@ -29,7 +28,7 @@ const createAssistant = async (req, res) => {
|
|||
delete assistantData.conversation_starters;
|
||||
delete assistantData.append_current_datetime;
|
||||
|
||||
const toolDefinitions = (await getCachedTools()) ?? {};
|
||||
const toolDefinitions = await getAssistantToolDefinitions({ req, tools });
|
||||
const healedTools = await healMcpToolNames({ req, tools, toolDefinitions });
|
||||
|
||||
assistantData.tools = healedTools
|
||||
|
|
@ -135,7 +134,7 @@ const updateAssistant = async ({ req, openai, assistant_id, updateData }) => {
|
|||
}
|
||||
|
||||
let hasFileSearch = false;
|
||||
const toolDefinitions = (await getCachedTools()) ?? {};
|
||||
const toolDefinitions = await getAssistantToolDefinitions({ req, tools: updateData.tools });
|
||||
const healedTools = await healMcpToolNames({ req, tools: updateData.tools, toolDefinitions });
|
||||
for (const tool of healedTools) {
|
||||
/** Agents-runtime-only tools (e.g. ask_user_question) cannot execute on
|
||||
|
|
|
|||
|
|
@ -35,7 +35,12 @@ const {
|
|||
resolveMcpConfigNames,
|
||||
resolveAllMcpConfigs,
|
||||
} = require('~/server/services/MCP');
|
||||
const { cacheMCPServerTools, getMCPServerTools } = require('~/server/services/Config');
|
||||
const {
|
||||
cacheMCPServerTools,
|
||||
getMCPServerTools,
|
||||
getMCPToolsCacheGeneration,
|
||||
invalidateCachedTools,
|
||||
} = require('~/server/services/Config');
|
||||
const { getResourcePermissionsMap } = require('~/server/services/PermissionService');
|
||||
const { hasCapability } = require('~/server/middleware/roles/capabilities');
|
||||
const { getMCPManager, getMCPServersRegistry } = require('~/config');
|
||||
|
|
@ -94,6 +99,68 @@ function handleMCPError(error, res) {
|
|||
return null;
|
||||
}
|
||||
|
||||
/** Disposes a stale local connection after its DB-backed config has changed. */
|
||||
async function disconnectLocalMCPServer(userId, serverName) {
|
||||
try {
|
||||
await getMCPManager()?.disconnectUserConnection(userId, serverName);
|
||||
} catch (error) {
|
||||
logger.warn(
|
||||
`[MCP Cache] Failed to disconnect the local connection for ${serverName} (user: ${userId}):`,
|
||||
error,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const POST_COMMIT_FENCE_RETRY_DELAYS_MS = [0, 50, 200];
|
||||
|
||||
/** Retries the shared fence after persistence; config-bound connections remain a durable
|
||||
* fallback if Redis stays unavailable, so an old connection cannot serve the new config. */
|
||||
async function fenceCommittedMCPMutation({ userId, serverName }) {
|
||||
let lastError;
|
||||
for (const delay of POST_COMMIT_FENCE_RETRY_DELAYS_MS) {
|
||||
if (delay > 0) {
|
||||
await new Promise((resolve) => setTimeout(resolve, delay));
|
||||
}
|
||||
try {
|
||||
await invalidateCachedTools({ userId, serverName });
|
||||
return;
|
||||
} catch (error) {
|
||||
lastError = error;
|
||||
logger.warn(
|
||||
`[MCP Cache] Failed to fence committed mutation for ${serverName} (user: ${userId}); retrying:`,
|
||||
error,
|
||||
);
|
||||
}
|
||||
}
|
||||
throw lastError;
|
||||
}
|
||||
|
||||
/**
|
||||
* Republishes the pre-mutation catalog under the new fence when persistence
|
||||
* fails. The retained connection will reacquire that generation on its next
|
||||
* use; this snapshot keeps every replica authoritative in the meantime.
|
||||
*/
|
||||
async function restoreRetainedServerCatalog({ userId, serverName, serverConfig, serverTools }) {
|
||||
if (serverTools == null) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const publicationGeneration = await getMCPToolsCacheGeneration({ userId, serverName });
|
||||
await cacheMCPServerTools({
|
||||
userId,
|
||||
serverName,
|
||||
serverConfig,
|
||||
serverTools,
|
||||
publicationGeneration,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error(
|
||||
`[MCP Cache] Failed to restore the retained catalog for ${serverName} (user: ${userId}):`,
|
||||
error,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all MCP tools available to the user.
|
||||
*/
|
||||
|
|
@ -150,8 +217,14 @@ const getMCPTools = async (req, res) => {
|
|||
}
|
||||
|
||||
let serverTools;
|
||||
let publicationGeneration;
|
||||
try {
|
||||
serverTools = await mcpManager.getServerToolFunctions(userId, serverName);
|
||||
({ tools: serverTools, publicationGeneration } =
|
||||
await mcpManager.getServerToolFunctionsSnapshot(
|
||||
userId,
|
||||
serverName,
|
||||
mcpConfig[serverName],
|
||||
));
|
||||
} catch (error) {
|
||||
logger.error(`[getMCPTools] Error fetching tools for server ${serverName}:`, error);
|
||||
continue;
|
||||
|
|
@ -162,17 +235,16 @@ const getMCPTools = async (req, res) => {
|
|||
}
|
||||
serverToolsMap.set(serverName, serverTools);
|
||||
|
||||
if (Object.keys(serverTools).length > 0) {
|
||||
// Cache asynchronously without blocking
|
||||
cacheMCPServerTools({
|
||||
userId,
|
||||
serverName,
|
||||
serverTools,
|
||||
serverConfig: mcpConfig[serverName],
|
||||
}).catch((err) =>
|
||||
logger.error(`[getMCPTools] Failed to cache tools for ${serverName}:`, err),
|
||||
);
|
||||
}
|
||||
// Empty is an authoritative catalog too; re-cache it after TTL expiry to avoid polling.
|
||||
cacheMCPServerTools({
|
||||
userId,
|
||||
serverName,
|
||||
serverTools,
|
||||
serverConfig: mcpConfig[serverName],
|
||||
publicationGeneration,
|
||||
}).catch((err) =>
|
||||
logger.error(`[getMCPTools] Failed to cache tools for ${serverName}:`, err),
|
||||
);
|
||||
}
|
||||
|
||||
// Process each configured server
|
||||
|
|
@ -509,12 +581,30 @@ const updateMCPServerController = async (req, res) => {
|
|||
.json({ message: 'Forbidden: Insufficient permissions to configure OBO' });
|
||||
}
|
||||
|
||||
const parsedConfig = await getMCPServersRegistry().updateServer(
|
||||
const registry = getMCPServersRegistry();
|
||||
const parsedConfig = await registry.inspectServerUpdate(
|
||||
serverName,
|
||||
validation.data,
|
||||
'DB',
|
||||
userId,
|
||||
);
|
||||
const retainedTools = await getMCPServerTools(userId, serverName, existingConfig);
|
||||
await invalidateCachedTools({ userId, serverName });
|
||||
try {
|
||||
await registry.commitServerUpdate(serverName, parsedConfig, 'DB', userId);
|
||||
} catch (error) {
|
||||
await restoreRetainedServerCatalog({
|
||||
userId,
|
||||
serverName,
|
||||
serverConfig: existingConfig,
|
||||
serverTools: retainedTools,
|
||||
});
|
||||
throw error;
|
||||
}
|
||||
/** Fence connections another replica could have created from the old DB
|
||||
* config between the pre-commit fence and the committed update. */
|
||||
await fenceCommittedMCPMutation({ userId, serverName });
|
||||
await disconnectLocalMCPServer(userId, serverName);
|
||||
|
||||
res.status(200).json(redactServerSecrets(parsedConfig, { canEdit: true }));
|
||||
} catch (error) {
|
||||
|
|
@ -535,7 +625,24 @@ const deleteMCPServerController = async (req, res) => {
|
|||
try {
|
||||
const userId = req.user?.id;
|
||||
const { serverName } = req.params;
|
||||
await getMCPServersRegistry().removeServer(serverName, 'DB', userId);
|
||||
const registry = getMCPServersRegistry();
|
||||
const existingConfig = await registry.getServerConfig(serverName, userId);
|
||||
const retainedTools = await getMCPServerTools(userId, serverName, existingConfig);
|
||||
await invalidateCachedTools({ userId, serverName });
|
||||
try {
|
||||
await registry.removeServer(serverName, 'DB', userId);
|
||||
} catch (error) {
|
||||
await restoreRetainedServerCatalog({
|
||||
userId,
|
||||
serverName,
|
||||
serverConfig: existingConfig,
|
||||
serverTools: retainedTools,
|
||||
});
|
||||
throw error;
|
||||
}
|
||||
/** Fence connections another replica could have created before deletion committed. */
|
||||
await fenceCommittedMCPMutation({ userId, serverName });
|
||||
await disconnectLocalMCPServer(userId, serverName);
|
||||
res.status(200).json({ message: 'MCP server deleted successfully' });
|
||||
} catch (error) {
|
||||
logger.error('[deleteMCPServer]', error);
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ jest.mock('~/server/services/Config', () => ({
|
|||
fileStrategy: 'local',
|
||||
imageOutputType: 'PNG',
|
||||
}),
|
||||
mergeAppTools: jest.fn().mockResolvedValue(undefined),
|
||||
setCachedTools: jest.fn(),
|
||||
}));
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ const { MongoMemoryServer } = require('mongodb-memory-server');
|
|||
const mongoose = require('mongoose');
|
||||
|
||||
jest.mock('~/server/services/Config', () => ({
|
||||
mergeAppTools: jest.fn().mockResolvedValue(undefined),
|
||||
loadCustomConfig: jest.fn(() => Promise.resolve({})),
|
||||
getAppConfig: jest.fn().mockResolvedValue({
|
||||
paths: {
|
||||
|
|
|
|||
|
|
@ -20,6 +20,8 @@ const mockRegistryInstance = {
|
|||
getAllServerConfigs: jest.fn(),
|
||||
ensureConfigServers: jest.fn().mockResolvedValue({}),
|
||||
addServer: jest.fn(),
|
||||
inspectServerUpdate: jest.fn(),
|
||||
commitServerUpdate: jest.fn(),
|
||||
updateServer: jest.fn(),
|
||||
removeServer: jest.fn(),
|
||||
getAllowedDomains: jest.fn().mockReturnValue(null),
|
||||
|
|
@ -128,6 +130,9 @@ jest.mock('~/models', () => ({
|
|||
jest.mock('~/server/services/Config', () => ({
|
||||
setCachedTools: jest.fn(),
|
||||
getCachedTools: jest.fn(),
|
||||
cacheMCPServerTools: jest.fn(),
|
||||
getMCPToolsCacheGeneration: jest.fn().mockResolvedValue('test-generation'),
|
||||
invalidateCachedTools: jest.fn(),
|
||||
getMCPServerTools: jest.fn(),
|
||||
loadCustomConfig: jest.fn(),
|
||||
getAppConfig: jest.fn().mockResolvedValue({ mcpConfig: {} }),
|
||||
|
|
@ -272,8 +277,15 @@ describe('MCP Routes', () => {
|
|||
*/
|
||||
mockRegistryInstance.getServerConfig.mockReset().mockResolvedValue(undefined);
|
||||
mockRegistryInstance.addServer.mockReset();
|
||||
mockRegistryInstance.inspectServerUpdate.mockReset();
|
||||
mockRegistryInstance.commitServerUpdate.mockReset();
|
||||
mockRegistryInstance.updateServer.mockReset();
|
||||
mockRegistryInstance.removeServer.mockReset();
|
||||
const cacheService = require('~/server/services/Config');
|
||||
cacheService.getMCPServerTools.mockReset().mockResolvedValue(null);
|
||||
cacheService.getMCPToolsCacheGeneration.mockReset().mockResolvedValue('test-generation');
|
||||
cacheService.cacheMCPServerTools.mockReset().mockResolvedValue(undefined);
|
||||
cacheService.invalidateCachedTools.mockReset().mockResolvedValue(undefined);
|
||||
});
|
||||
|
||||
describe('GET /:serverName/oauth/initiate', () => {
|
||||
|
|
@ -843,7 +855,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
const mockMcpManager = {
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchTools: jest.fn().mockResolvedValue([]),
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
|
||||
}),
|
||||
};
|
||||
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
|
||||
|
|
@ -905,7 +917,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
const mockMcpManager = {
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchTools: jest.fn().mockResolvedValue([]),
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
|
||||
}),
|
||||
};
|
||||
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
|
||||
|
|
@ -958,9 +970,12 @@ describe('MCP Routes', () => {
|
|||
mockRegistryInstance.getServerConfig.mockResolvedValue({});
|
||||
mockResolveAllMcpConfigs.mockResolvedValueOnce({ 'test-server': mergedServerConfig });
|
||||
|
||||
const fetchOrderedToolsSnapshot = jest
|
||||
.fn()
|
||||
.mockResolvedValue({ tools: fetchedTools, complete: true });
|
||||
const mockMcpManager = {
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchTools: jest.fn().mockResolvedValue(fetchedTools),
|
||||
fetchOrderedToolsSnapshot,
|
||||
}),
|
||||
};
|
||||
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
|
||||
|
|
@ -976,6 +991,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
expect(response.status).toBe(302);
|
||||
expect(mockResolveAllMcpConfigs).toHaveBeenCalledWith('test-user-id');
|
||||
expect(fetchOrderedToolsSnapshot).toHaveBeenCalledTimes(1);
|
||||
expect(mockMcpManager.getUserConnection).toHaveBeenCalledWith(
|
||||
expect.objectContaining({ serverConfig: mergedServerConfig }),
|
||||
);
|
||||
|
|
@ -1030,7 +1046,9 @@ describe('MCP Routes', () => {
|
|||
|
||||
const mockMcpManager = {
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchTools: jest.fn().mockResolvedValue(fetchedTools),
|
||||
fetchToolsSnapshot: jest
|
||||
.fn()
|
||||
.mockResolvedValue({ tools: fetchedTools, complete: true }),
|
||||
}),
|
||||
};
|
||||
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
|
||||
|
|
@ -1093,7 +1111,9 @@ describe('MCP Routes', () => {
|
|||
|
||||
const mockMcpManager = {
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchTools: jest.fn().mockResolvedValue(fetchedTools),
|
||||
fetchToolsSnapshot: jest
|
||||
.fn()
|
||||
.mockResolvedValue({ tools: fetchedTools, complete: true }),
|
||||
}),
|
||||
};
|
||||
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
|
||||
|
|
@ -1209,13 +1229,16 @@ describe('MCP Routes', () => {
|
|||
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
|
||||
|
||||
const mockUserConnection = {
|
||||
fetchTools: jest.fn().mockResolvedValue([
|
||||
{
|
||||
name: 'test-tool',
|
||||
description: 'A test tool',
|
||||
inputSchema: { type: 'object' },
|
||||
},
|
||||
]),
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({
|
||||
tools: [
|
||||
{
|
||||
name: 'test-tool',
|
||||
description: 'A test tool',
|
||||
inputSchema: { type: 'object' },
|
||||
},
|
||||
],
|
||||
complete: true,
|
||||
}),
|
||||
};
|
||||
const mockMcpManager = {
|
||||
getUserConnection: jest.fn().mockResolvedValue(mockUserConnection),
|
||||
|
|
@ -1310,7 +1333,7 @@ describe('MCP Routes', () => {
|
|||
});
|
||||
require('~/config').getMCPManager.mockReturnValue({
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchTools: jest.fn().mockResolvedValue([]),
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
|
||||
}),
|
||||
});
|
||||
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
|
||||
|
|
@ -1385,7 +1408,7 @@ describe('MCP Routes', () => {
|
|||
});
|
||||
require('~/config').getMCPManager.mockReturnValue({
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchTools: jest.fn().mockResolvedValue([]),
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
|
||||
}),
|
||||
});
|
||||
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
|
||||
|
|
@ -1443,7 +1466,7 @@ describe('MCP Routes', () => {
|
|||
});
|
||||
require('~/config').getMCPManager.mockReturnValue({
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchTools: jest.fn().mockResolvedValue([]),
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
|
||||
}),
|
||||
});
|
||||
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
|
||||
|
|
@ -1497,7 +1520,7 @@ describe('MCP Routes', () => {
|
|||
});
|
||||
require('~/config').getMCPManager.mockReturnValue({
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchTools: jest.fn().mockResolvedValue([]),
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
|
||||
}),
|
||||
});
|
||||
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
|
||||
|
|
@ -1730,7 +1753,7 @@ describe('MCP Routes', () => {
|
|||
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
|
||||
|
||||
const mockUserConnection = {
|
||||
fetchTools: jest.fn().mockResolvedValue([]),
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
|
||||
};
|
||||
const mockMcpManager = {
|
||||
getUserConnection: jest.fn().mockResolvedValue(mockUserConnection),
|
||||
|
|
@ -2389,10 +2412,13 @@ describe('MCP Routes', () => {
|
|||
|
||||
it('should successfully reinitialize server and cache tools', async () => {
|
||||
const mockUserConnection = {
|
||||
fetchTools: jest.fn().mockResolvedValue([
|
||||
{ name: 'tool1', description: 'Test tool 1', inputSchema: { type: 'object' } },
|
||||
{ name: 'tool2', description: 'Test tool 2', inputSchema: { type: 'object' } },
|
||||
]),
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({
|
||||
tools: [
|
||||
{ name: 'tool1', description: 'Test tool 1', inputSchema: { type: 'object' } },
|
||||
{ name: 'tool2', description: 'Test tool 2', inputSchema: { type: 'object' } },
|
||||
],
|
||||
complete: true,
|
||||
}),
|
||||
};
|
||||
|
||||
const mockMcpManager = {
|
||||
|
|
@ -2435,11 +2461,18 @@ describe('MCP Routes', () => {
|
|||
'test-user-id',
|
||||
'test-server',
|
||||
);
|
||||
expect(require('~/server/services/Config').invalidateCachedTools).toHaveBeenCalledWith({
|
||||
userId: 'test-user-id',
|
||||
serverName: 'test-server',
|
||||
});
|
||||
expect(
|
||||
require('~/server/services/Config').invalidateCachedTools.mock.invocationCallOrder[0],
|
||||
).toBeLessThan(mockMcpManager.disconnectUserConnection.mock.invocationCallOrder[0]);
|
||||
});
|
||||
|
||||
it('should handle server with custom user variables', async () => {
|
||||
const mockUserConnection = {
|
||||
fetchTools: jest.fn().mockResolvedValue([]),
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
|
||||
};
|
||||
|
||||
const mockMcpManager = {
|
||||
|
|
@ -2896,7 +2929,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
const mockMcpManager = {
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchTools: jest.fn().mockResolvedValue([]),
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }),
|
||||
}),
|
||||
};
|
||||
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
|
||||
|
|
@ -2949,10 +2982,12 @@ describe('MCP Routes', () => {
|
|||
|
||||
const mockMcpManager = {
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchTools: jest
|
||||
.fn()
|
||||
.mockResolvedValue([{ name: 'test-tool', description: 'Test tool' }]),
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({
|
||||
tools: [{ name: 'test-tool', description: 'Test tool' }],
|
||||
complete: true,
|
||||
}),
|
||||
}),
|
||||
getToolPublicationGeneration: jest.fn().mockReturnValue('oauth-connection-generation'),
|
||||
};
|
||||
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
|
||||
|
||||
|
|
@ -2967,6 +3002,60 @@ describe('MCP Routes', () => {
|
|||
const basePath = getBasePath();
|
||||
|
||||
expect(response.headers.location).toContain(`${basePath}/oauth/success`);
|
||||
expect(require('~/server/services/Config/mcp').updateMCPServerTools).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
userId: 'test-user-id',
|
||||
serverName: 'test-server',
|
||||
tools: [{ name: 'test-tool', description: 'Test tool' }],
|
||||
publicationGeneration: 'oauth-connection-generation',
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('preserves cached tools when the post-OAuth snapshot is incomplete', async () => {
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { MCPOAuthHandler, MCPTokenStorage } = require('@librechat/api');
|
||||
const mockTokens = {
|
||||
access_token: 'edge-access-token',
|
||||
refresh_token: 'edge-refresh-token',
|
||||
};
|
||||
const mockFlowManager = {
|
||||
getFlowState: jest.fn(),
|
||||
completeFlow: jest.fn(),
|
||||
};
|
||||
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
|
||||
MCPOAuthHandler.getFlowState.mockResolvedValue({
|
||||
state: 'test-user-id:test-server',
|
||||
serverName: 'test-server',
|
||||
userId: 'test-user-id',
|
||||
metadata: { serverUrl: 'https://example.com', oauth: {} },
|
||||
clientInfo: {},
|
||||
codeVerifier: 'test-verifier',
|
||||
});
|
||||
mockOAuthCompletion(mockTokens);
|
||||
MCPTokenStorage.storeTokens.mockResolvedValue();
|
||||
mockRegistryInstance.getServerConfig.mockResolvedValue({});
|
||||
require('~/config').getMCPManager.mockReturnValue({
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchToolsSnapshot: jest.fn().mockResolvedValue({
|
||||
tools: [{ name: 'partial-tool', description: 'Only the first page' }],
|
||||
complete: false,
|
||||
}),
|
||||
}),
|
||||
getToolPublicationGeneration: jest.fn().mockReturnValue('oauth-connection-generation'),
|
||||
});
|
||||
|
||||
const flowId = 'test-user-id:test-server';
|
||||
const csrfToken = generateTestCsrfToken(flowId);
|
||||
await request(app)
|
||||
.get(`/api/mcp/test-server/oauth/callback?code=test-code&state=${flowId}`)
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.expect(302);
|
||||
|
||||
expect(require('~/server/services/Config/mcp').updateMCPServerTools).not.toHaveBeenCalled();
|
||||
expect(logger.warn).toHaveBeenCalledWith(
|
||||
'[MCP OAuth] Preserving cached tools for test-server because tools/list returned an incomplete snapshot',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -3064,6 +3153,68 @@ describe('MCP Routes', () => {
|
|||
expect(mockResolveAllMcpConfigs).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('caches a live user snapshot with its connection-bound publication generation', async () => {
|
||||
const { Constants } = require('librechat-data-provider');
|
||||
const { cacheMCPServerTools, getMCPServerTools } = require('~/server/services/Config');
|
||||
const pluginKey = `search${Constants.mcp_delimiter}user-server`;
|
||||
const serverTools = {
|
||||
[pluginKey]: {
|
||||
type: 'function',
|
||||
function: {
|
||||
name: pluginKey,
|
||||
description: 'Search',
|
||||
parameters: { type: 'object' },
|
||||
},
|
||||
},
|
||||
};
|
||||
const serverConfig = { type: 'sse', url: 'https://user.example.com/sse' };
|
||||
mockResolveAllMcpConfigs.mockResolvedValueOnce({ 'user-server': serverConfig });
|
||||
getMCPServerTools.mockResolvedValueOnce(null);
|
||||
cacheMCPServerTools.mockResolvedValueOnce();
|
||||
require('~/config').getMCPManager.mockReturnValue({
|
||||
getServerToolFunctionsSnapshot: jest.fn().mockResolvedValue({
|
||||
tools: serverTools,
|
||||
publicationGeneration: 'connection-generation',
|
||||
}),
|
||||
});
|
||||
|
||||
const response = await request(app).get('/api/mcp/tools');
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(cacheMCPServerTools).toHaveBeenCalledWith({
|
||||
userId: 'test-user-id',
|
||||
serverName: 'user-server',
|
||||
serverTools,
|
||||
serverConfig,
|
||||
publicationGeneration: 'connection-generation',
|
||||
});
|
||||
});
|
||||
|
||||
it('re-caches an authoritative empty live snapshot after a cache miss', async () => {
|
||||
const { cacheMCPServerTools, getMCPServerTools } = require('~/server/services/Config');
|
||||
const serverConfig = { type: 'sse', url: 'https://empty.example.com/sse' };
|
||||
mockResolveAllMcpConfigs.mockResolvedValueOnce({ 'empty-server': serverConfig });
|
||||
getMCPServerTools.mockResolvedValueOnce(null);
|
||||
cacheMCPServerTools.mockResolvedValueOnce();
|
||||
const getServerToolFunctionsSnapshot = jest.fn().mockResolvedValue({
|
||||
tools: {},
|
||||
publicationGeneration: undefined,
|
||||
});
|
||||
require('~/config').getMCPManager.mockReturnValue({ getServerToolFunctionsSnapshot });
|
||||
|
||||
const response = await request(app).get('/api/mcp/tools');
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body.servers['empty-server'].tools).toEqual([]);
|
||||
expect(cacheMCPServerTools).toHaveBeenCalledWith({
|
||||
userId: 'test-user-id',
|
||||
serverName: 'empty-server',
|
||||
serverTools: {},
|
||||
serverConfig,
|
||||
publicationGeneration: undefined,
|
||||
});
|
||||
});
|
||||
|
||||
it('should continue returning MCP tools when one server cache lookup fails', async () => {
|
||||
const { Constants } = require('librechat-data-provider');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
|
|
@ -3095,9 +3246,12 @@ describe('MCP Routes', () => {
|
|||
},
|
||||
});
|
||||
|
||||
const mockGetServerToolFunctions = jest.fn().mockResolvedValue(null);
|
||||
const mockGetServerToolFunctionsSnapshot = jest.fn().mockResolvedValue({
|
||||
tools: null,
|
||||
publicationGeneration: 'test-generation',
|
||||
});
|
||||
require('~/config').getMCPManager.mockReturnValue({
|
||||
getServerToolFunctions: mockGetServerToolFunctions,
|
||||
getServerToolFunctionsSnapshot: mockGetServerToolFunctionsSnapshot,
|
||||
});
|
||||
|
||||
const response = await request(app).get('/api/mcp/tools');
|
||||
|
|
@ -3107,7 +3261,14 @@ describe('MCP Routes', () => {
|
|||
'[getMCPTools] Error fetching cached tools for bad-server:',
|
||||
expect.any(Error),
|
||||
);
|
||||
expect(mockGetServerToolFunctions).toHaveBeenCalledWith('test-user-id', 'bad-server');
|
||||
expect(mockGetServerToolFunctionsSnapshot).toHaveBeenCalledWith(
|
||||
'test-user-id',
|
||||
'bad-server',
|
||||
{
|
||||
type: 'sse',
|
||||
url: 'https://bad.example.com/sse',
|
||||
},
|
||||
);
|
||||
expect(response.body.servers['good-server']).toMatchObject({
|
||||
name: 'good-server',
|
||||
icon: '/icons/good.svg',
|
||||
|
|
@ -3142,9 +3303,12 @@ describe('MCP Routes', () => {
|
|||
|
||||
getMCPServerTools.mockRejectedValue(new Error('cache unavailable'));
|
||||
|
||||
const mockGetServerToolFunctions = jest.fn().mockResolvedValue(null);
|
||||
const mockGetServerToolFunctionsSnapshot = jest.fn().mockResolvedValue({
|
||||
tools: null,
|
||||
publicationGeneration: 'test-generation',
|
||||
});
|
||||
require('~/config').getMCPManager.mockReturnValue({
|
||||
getServerToolFunctions: mockGetServerToolFunctions,
|
||||
getServerToolFunctionsSnapshot: mockGetServerToolFunctionsSnapshot,
|
||||
});
|
||||
|
||||
const response = await request(app).get('/api/mcp/tools');
|
||||
|
|
@ -3159,7 +3323,7 @@ describe('MCP Routes', () => {
|
|||
tools: [],
|
||||
});
|
||||
expect(logger.error).toHaveBeenCalledTimes(2);
|
||||
expect(mockGetServerToolFunctions).toHaveBeenCalledTimes(2);
|
||||
expect(mockGetServerToolFunctionsSnapshot).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -3567,7 +3731,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
expect(response.status).toBe(403);
|
||||
expect(response.body.message).toMatch(/Insufficient permissions to configure OBO/);
|
||||
expect(mockRegistryInstance.updateServer).not.toHaveBeenCalled();
|
||||
expect(mockRegistryInstance.inspectServerUpdate).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('allows PATCH without CONFIGURE_OBO when OBO is unchanged', async () => {
|
||||
|
|
@ -3589,10 +3753,11 @@ describe('MCP Routes', () => {
|
|||
...oboConfig,
|
||||
obo: { scopes: 'api://mcp-server-id/Mcp.Tools.ReadWrite' },
|
||||
});
|
||||
mockRegistryInstance.updateServer.mockResolvedValue({
|
||||
mockRegistryInstance.inspectServerUpdate.mockResolvedValue({
|
||||
...oboConfig,
|
||||
title: 'Renamed OBO Server',
|
||||
});
|
||||
mockRegistryInstance.commitServerUpdate.mockResolvedValue(undefined);
|
||||
|
||||
const response = await request(app)
|
||||
.patch('/api/mcp/servers/obo-server')
|
||||
|
|
@ -3605,7 +3770,7 @@ describe('MCP Routes', () => {
|
|||
});
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(mockRegistryInstance.updateServer).toHaveBeenCalled();
|
||||
expect(mockRegistryInstance.commitServerUpdate).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('rejects PATCH that removes OBO from an existing OBO server without CONFIGURE_OBO', async () => {
|
||||
|
|
@ -3639,7 +3804,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
expect(response.status).toBe(403);
|
||||
expect(response.body.message).toMatch(/Insufficient permissions to configure OBO/);
|
||||
expect(mockRegistryInstance.updateServer).not.toHaveBeenCalled();
|
||||
expect(mockRegistryInstance.inspectServerUpdate).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('rejects PATCH that redirects the URL of an existing OBO server without CONFIGURE_OBO', async () => {
|
||||
|
|
@ -3676,7 +3841,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
expect(response.status).toBe(403);
|
||||
expect(response.body.message).toMatch(/Insufficient permissions to configure OBO/);
|
||||
expect(mockRegistryInstance.updateServer).not.toHaveBeenCalled();
|
||||
expect(mockRegistryInstance.inspectServerUpdate).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -3772,7 +3937,11 @@ describe('MCP Routes', () => {
|
|||
description: 'Updated description',
|
||||
};
|
||||
|
||||
mockRegistryInstance.updateServer.mockResolvedValue({ ...updatedConfig, source: 'user' });
|
||||
mockRegistryInstance.inspectServerUpdate.mockResolvedValue({
|
||||
...updatedConfig,
|
||||
source: 'user',
|
||||
});
|
||||
mockRegistryInstance.commitServerUpdate.mockResolvedValue(undefined);
|
||||
|
||||
const response = await request(app)
|
||||
.patch('/api/mcp/servers/test-server')
|
||||
|
|
@ -3782,7 +3951,7 @@ describe('MCP Routes', () => {
|
|||
expect(response.body.type).toBe('sse');
|
||||
expect(response.body.url).toBe('https://updated-mcp-server.example.com/sse');
|
||||
expect(response.body.title).toBe('Updated Server');
|
||||
expect(mockRegistryInstance.updateServer).toHaveBeenCalledWith(
|
||||
expect(mockRegistryInstance.inspectServerUpdate).toHaveBeenCalledWith(
|
||||
'test-server',
|
||||
expect.objectContaining({
|
||||
type: 'sse',
|
||||
|
|
@ -3800,13 +3969,14 @@ describe('MCP Routes', () => {
|
|||
title: 'Updated Server',
|
||||
};
|
||||
|
||||
mockRegistryInstance.updateServer.mockResolvedValue({
|
||||
mockRegistryInstance.inspectServerUpdate.mockResolvedValue({
|
||||
...validConfig,
|
||||
apiKey: { source: 'admin', authorization_type: 'bearer', key: 'preserved-admin-key' },
|
||||
oauth: { client_id: 'cid', client_secret: 'preserved-oauth-secret' },
|
||||
headers: { Authorization: 'Bearer internal-token' },
|
||||
env: { DATABASE_URL: 'postgres://admin:pass@localhost/db' },
|
||||
});
|
||||
mockRegistryInstance.commitServerUpdate.mockResolvedValue(undefined);
|
||||
|
||||
const response = await request(app)
|
||||
.patch('/api/mcp/servers/test-server')
|
||||
|
|
@ -3832,7 +4002,7 @@ describe('MCP Routes', () => {
|
|||
statusCode: 400,
|
||||
},
|
||||
);
|
||||
mockRegistryInstance.updateServer.mockRejectedValue(error);
|
||||
mockRegistryInstance.inspectServerUpdate.mockRejectedValue(error);
|
||||
|
||||
const response = await request(app)
|
||||
.patch('/api/mcp/servers/test-server')
|
||||
|
|
@ -3884,7 +4054,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
expect(response.status).toBe(400);
|
||||
expect(response.body.message).toBe('Invalid configuration');
|
||||
expect(mockRegistryInstance.updateServer).not.toHaveBeenCalled();
|
||||
expect(mockRegistryInstance.inspectServerUpdate).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should reject streamable-http URL containing env variable references', async () => {
|
||||
|
|
@ -3899,7 +4069,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
expect(response.status).toBe(400);
|
||||
expect(response.body.message).toBe('Invalid configuration');
|
||||
expect(mockRegistryInstance.updateServer).not.toHaveBeenCalled();
|
||||
expect(mockRegistryInstance.inspectServerUpdate).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should reject websocket URL containing env variable references', async () => {
|
||||
|
|
@ -3914,7 +4084,7 @@ describe('MCP Routes', () => {
|
|||
|
||||
expect(response.status).toBe(400);
|
||||
expect(response.body.message).toBe('Invalid configuration');
|
||||
expect(mockRegistryInstance.updateServer).not.toHaveBeenCalled();
|
||||
expect(mockRegistryInstance.inspectServerUpdate).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should return 500 when registry throws error', async () => {
|
||||
|
|
@ -3924,7 +4094,7 @@ describe('MCP Routes', () => {
|
|||
title: 'Test Server',
|
||||
};
|
||||
|
||||
mockRegistryInstance.updateServer.mockRejectedValue(new Error('Update failed'));
|
||||
mockRegistryInstance.inspectServerUpdate.mockRejectedValue(new Error('Update failed'));
|
||||
|
||||
const response = await request(app)
|
||||
.patch('/api/mcp/servers/test-server')
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ const {
|
|||
} = require('~/server/services/MCP');
|
||||
const { requireJwtAuth, canAccessMCPServerResource } = require('~/server/middleware');
|
||||
const { getUserPluginAuthValue } = require('~/server/services/PluginService');
|
||||
const { invalidateCachedTools } = require('~/server/services/Config');
|
||||
const { updateMCPServerTools } = require('~/server/services/Config/mcp');
|
||||
const { reinitMCPServer } = require('~/server/services/Tools/mcp');
|
||||
const { getLogStores } = require('~/cache');
|
||||
|
|
@ -542,13 +543,24 @@ router.get('/:serverName/oauth/callback', async (req, res) => {
|
|||
const oauthReconnectionManager = getOAuthReconnectionManager();
|
||||
oauthReconnectionManager.clearReconnection(flowState.userId, serverName);
|
||||
|
||||
const tools = await userConnection.fetchTools();
|
||||
await updateMCPServerTools({
|
||||
userId: flowState.userId,
|
||||
serverName,
|
||||
tools,
|
||||
serverConfig,
|
||||
});
|
||||
const snapshot =
|
||||
typeof userConnection.fetchOrderedToolsSnapshot === 'function'
|
||||
? await userConnection.fetchOrderedToolsSnapshot()
|
||||
: await userConnection.fetchToolsSnapshot();
|
||||
if (snapshot.complete) {
|
||||
const publicationGeneration = mcpManager.getToolPublicationGeneration?.(userConnection);
|
||||
await updateMCPServerTools({
|
||||
userId: flowState.userId,
|
||||
serverName,
|
||||
tools: snapshot.tools,
|
||||
serverConfig,
|
||||
publicationGeneration,
|
||||
});
|
||||
} else {
|
||||
logger.warn(
|
||||
`[MCP OAuth] Preserving cached tools for ${serverName} because tools/list returned an incomplete snapshot`,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
logger.debug(`[MCP OAuth] System-level OAuth completed for ${serverName}`);
|
||||
}
|
||||
|
|
@ -791,7 +803,11 @@ router.post(
|
|||
});
|
||||
}
|
||||
|
||||
await mcpManager.disconnectUserConnection(user.id, serverName);
|
||||
try {
|
||||
await invalidateCachedTools({ userId: user.id, serverName });
|
||||
} finally {
|
||||
await mcpManager.disconnectUserConnection(user.id, serverName);
|
||||
}
|
||||
logger.info(
|
||||
`[MCP Reinitialize] Disconnected existing user connection for server: ${serverName}`,
|
||||
);
|
||||
|
|
|
|||
341
api/server/services/Config/__tests__/getCachedTools.lock.spec.js
Normal file
341
api/server/services/Config/__tests__/getCachedTools.lock.spec.js
Normal file
|
|
@ -0,0 +1,341 @@
|
|||
const { CacheKeys } = require('librechat-data-provider');
|
||||
const calculateSlot = require('cluster-key-slot');
|
||||
|
||||
const mockRedisClient = {
|
||||
set: jest.fn(),
|
||||
eval: jest.fn(),
|
||||
};
|
||||
const mockKeyvRedisClient = {
|
||||
eval: jest.fn(),
|
||||
};
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
...jest.requireActual('@librechat/api'),
|
||||
cacheConfig: { FORCED_IN_MEMORY_CACHE_NAMESPACES: [] },
|
||||
mcpConfig: { USER_CONNECTION_IDLE_TIMEOUT: 15 * 60 * 1000 },
|
||||
ioredisClient: mockRedisClient,
|
||||
keyvRedisClient: mockKeyvRedisClient,
|
||||
}));
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: { warn: jest.fn() },
|
||||
}));
|
||||
|
||||
jest.mock('~/cache/getLogStores', () => jest.fn());
|
||||
|
||||
const getLogStores = require('~/cache/getLogStores');
|
||||
const mockCache = { get: jest.fn(), set: jest.fn(), delete: jest.fn() };
|
||||
getLogStores.mockReturnValue(mockCache);
|
||||
|
||||
const {
|
||||
getCachedTools,
|
||||
updateCachedGlobalTools,
|
||||
getMCPToolsCacheGeneration,
|
||||
setCachedTools,
|
||||
setCachedToolsIfCurrent,
|
||||
runWithGlobalCacheLock,
|
||||
invalidateCachedTools,
|
||||
setCachedToolsWithinGlobalLock,
|
||||
getNextAppToolsPublicationRevision,
|
||||
setCachedAppServerTools,
|
||||
} = require('../getCachedTools');
|
||||
|
||||
describe('global tool cache write lock', () => {
|
||||
afterEach(() => {
|
||||
jest.useRealTimers();
|
||||
});
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockRedisClient.set.mockResolvedValue('OK');
|
||||
mockRedisClient.eval.mockResolvedValue(1);
|
||||
mockKeyvRedisClient.eval.mockResolvedValue(1);
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
mockCache.delete.mockResolvedValue(true);
|
||||
});
|
||||
|
||||
it('acquires and safely releases the Redis lock around an aggregate update', async () => {
|
||||
const operation = jest.fn().mockResolvedValue('updated');
|
||||
|
||||
await expect(runWithGlobalCacheLock(operation)).resolves.toBe('updated');
|
||||
|
||||
expect(mockRedisClient.set).toHaveBeenCalledWith(
|
||||
`${CacheKeys.TOOL_CACHE}:tools:global:write-lock`,
|
||||
expect.any(String),
|
||||
'PX',
|
||||
30_000,
|
||||
'NX',
|
||||
);
|
||||
const token = mockRedisClient.set.mock.calls[0][1];
|
||||
expect(mockRedisClient.eval).toHaveBeenCalledWith(
|
||||
expect.stringContaining("redis.call('GET'"),
|
||||
1,
|
||||
`${CacheKeys.TOOL_CACHE}:tools:global:write-lock`,
|
||||
token,
|
||||
);
|
||||
const fenceKey = mockKeyvRedisClient.eval.mock.calls[0][1].keys[0];
|
||||
expect(calculateSlot(fenceKey)).toBe(calculateSlot(`${CacheKeys.TOOL_CACHE}:tools:global`));
|
||||
});
|
||||
|
||||
it('releases the Redis lock when the aggregate update fails', async () => {
|
||||
const operation = jest.fn().mockRejectedValue(new Error('cache read failed'));
|
||||
|
||||
await expect(runWithGlobalCacheLock(operation)).rejects.toThrow('cache read failed');
|
||||
|
||||
expect(mockRedisClient.eval).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('serializes direct global writes and invalidation', async () => {
|
||||
await setCachedTools({ builtin: {} });
|
||||
await invalidateCachedTools({ invalidateGlobal: true });
|
||||
|
||||
expect(mockRedisClient.set).toHaveBeenCalledTimes(2);
|
||||
expect(mockRedisClient.eval).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('does not reacquire the lock for a write already inside an aggregate update', async () => {
|
||||
await runWithGlobalCacheLock(() => setCachedToolsWithinGlobalLock({ mcp: {} }));
|
||||
|
||||
expect(mockRedisClient.set).toHaveBeenCalledTimes(1);
|
||||
expect(mockRedisClient.eval).toHaveBeenCalledTimes(1);
|
||||
expect(mockKeyvRedisClient.eval).toHaveBeenCalledWith(
|
||||
expect.stringContaining("redis.call('GET', KEYS[1])"),
|
||||
expect.objectContaining({
|
||||
keys: [
|
||||
`tools:global:write-fence:{${CacheKeys.TOOL_CACHE}:tools:global}`,
|
||||
`${CacheKeys.TOOL_CACHE}:tools:global`,
|
||||
],
|
||||
}),
|
||||
);
|
||||
expect(mockCache.set).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('atomically replaces a legacy global catalog while holding its Redis fence', async () => {
|
||||
mockCache.get.mockResolvedValue({ old_mcp_server: {}, builtin: {} });
|
||||
|
||||
await updateCachedGlobalTools(() => ({ builtin: {} }));
|
||||
|
||||
expect(mockRedisClient.set).toHaveBeenCalledTimes(1);
|
||||
expect(mockKeyvRedisClient.eval).toHaveBeenCalledTimes(3);
|
||||
expect(mockKeyvRedisClient.eval.mock.calls).toEqual(
|
||||
expect.arrayContaining([
|
||||
[
|
||||
expect.stringContaining("redis.call('PSETEX', KEYS[2]"),
|
||||
expect.objectContaining({
|
||||
keys: [
|
||||
`tools:global:write-fence:{${CacheKeys.TOOL_CACHE}:tools:global}`,
|
||||
`${CacheKeys.TOOL_CACHE}:tools:global`,
|
||||
],
|
||||
arguments: [
|
||||
expect.any(String),
|
||||
JSON.stringify({ builtin: {} }),
|
||||
expect.any(String),
|
||||
expect.any(String),
|
||||
],
|
||||
}),
|
||||
],
|
||||
]),
|
||||
);
|
||||
expect(mockCache.set).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('rejects a Redis-backed global write made without distributed lock ownership', async () => {
|
||||
await expect(setCachedToolsWithinGlobalLock({ unsafe: {} })).rejects.toThrow(
|
||||
'Global tool cache write requires lock ownership',
|
||||
);
|
||||
|
||||
expect(mockCache.set).not.toHaveBeenCalled();
|
||||
expect(mockKeyvRedisClient.eval).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('rejects a delayed global write after its distributed lease is lost', async () => {
|
||||
mockKeyvRedisClient.eval.mockResolvedValueOnce(1).mockResolvedValueOnce(0).mockResolvedValue(1);
|
||||
|
||||
await expect(
|
||||
runWithGlobalCacheLock(() => setCachedToolsWithinGlobalLock({ stale: {} })),
|
||||
).rejects.toThrow('Global tool cache lock ownership was lost before write');
|
||||
|
||||
expect(mockCache.set).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('rejects a delayed ownership claim after a newer owner has fenced the slot', async () => {
|
||||
const operation = jest.fn();
|
||||
mockKeyvRedisClient.eval.mockResolvedValueOnce(0);
|
||||
|
||||
await expect(runWithGlobalCacheLock(operation)).rejects.toThrow(
|
||||
'Tool cache lock expired or was superseded before ownership could be fenced',
|
||||
);
|
||||
|
||||
expect(operation).not.toHaveBeenCalled();
|
||||
const [claimScript, claimOptions] = mockKeyvRedisClient.eval.mock.calls[0];
|
||||
expect(claimScript).toContain('current ~= ARGV[1]');
|
||||
expect(claimScript).toContain("redis.call('TIME')");
|
||||
expect(claimOptions.arguments).toEqual([
|
||||
mockRedisClient.set.mock.calls[0][1],
|
||||
expect.any(String),
|
||||
'1000',
|
||||
]);
|
||||
});
|
||||
|
||||
it('atomically checks the generation and writes a generation-guarded user catalog', async () => {
|
||||
await expect(
|
||||
setCachedToolsIfCurrent(
|
||||
{ current: {} },
|
||||
{
|
||||
userId: 'user-1',
|
||||
serverName: 'server-1',
|
||||
configGeneration: 'config-current',
|
||||
publicationGeneration: 'generation-current',
|
||||
},
|
||||
),
|
||||
).resolves.toBe(true);
|
||||
|
||||
expect(mockRedisClient.set).toHaveBeenCalledWith(
|
||||
`${CacheKeys.TOOL_CACHE}:tools:mcp-write-lock:user-1:server-1`,
|
||||
expect.any(String),
|
||||
'PX',
|
||||
30_000,
|
||||
'NX',
|
||||
);
|
||||
expect(mockRedisClient.eval).toHaveBeenCalledTimes(1);
|
||||
expect(mockKeyvRedisClient.eval).toHaveBeenCalledWith(
|
||||
expect.stringContaining("redis.call('PSETEX', KEYS[2]"),
|
||||
expect.objectContaining({
|
||||
keys: [
|
||||
`${CacheKeys.TOOL_CACHE}:tools:metadata:mcp:user-generation:{user-1:server-1}`,
|
||||
`${CacheKeys.TOOL_CACHE}:tools:mcp:user:{user-1:server-1}:config-current`,
|
||||
],
|
||||
arguments: [
|
||||
'generation-current',
|
||||
expect.any(String),
|
||||
expect.any(String),
|
||||
expect.stringContaining('"publicationGeneration":"generation-current"'),
|
||||
expect.any(String),
|
||||
expect.any(String),
|
||||
],
|
||||
}),
|
||||
);
|
||||
expect(mockCache.set).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('does not write tools when the atomic generation check observes a replacement', async () => {
|
||||
mockKeyvRedisClient.eval.mockResolvedValue(0);
|
||||
|
||||
await expect(
|
||||
setCachedToolsIfCurrent(
|
||||
{ stale: {} },
|
||||
{
|
||||
userId: 'user-1',
|
||||
serverName: 'server-1',
|
||||
configGeneration: 'config-old',
|
||||
publicationGeneration: 'generation-old',
|
||||
},
|
||||
),
|
||||
).resolves.toBe(false);
|
||||
|
||||
expect(mockCache.set).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('orders app snapshots atomically in the app catalog Redis slot', async () => {
|
||||
mockCache.get.mockResolvedValue(null);
|
||||
mockKeyvRedisClient.eval
|
||||
.mockResolvedValueOnce(1)
|
||||
.mockResolvedValueOnce(2)
|
||||
.mockResolvedValueOnce(1)
|
||||
.mockResolvedValueOnce(0);
|
||||
|
||||
const older = await getNextAppToolsPublicationRevision('server-1', 'config-current');
|
||||
const newer = await getNextAppToolsPublicationRevision('server-1', 'config-current');
|
||||
await expect(
|
||||
setCachedAppServerTools('server-1', 'config-current', { current: {} }, newer),
|
||||
).resolves.toBe(true);
|
||||
await expect(
|
||||
setCachedAppServerTools('server-1', 'config-current', { stale: {} }, older),
|
||||
).resolves.toBe(false);
|
||||
|
||||
const [reserveScript, reserveOptions] = mockKeyvRedisClient.eval.mock.calls[0];
|
||||
const [writeScript, writeOptions] = mockKeyvRedisClient.eval.mock.calls[2];
|
||||
expect(reserveScript).toContain("redis.call('INCR', KEYS[1])");
|
||||
expect(writeScript).toContain('tonumber(current) > tonumber(ARGV[1])');
|
||||
expect(writeScript).toContain("currentEntry['value']['publicationRevision']");
|
||||
expect(calculateSlot(reserveOptions.keys[0])).toBe(calculateSlot(writeOptions.keys[1]));
|
||||
expect(calculateSlot(writeOptions.keys[0])).toBe(calculateSlot(writeOptions.keys[1]));
|
||||
expect(writeOptions.keys[0]).toContain('app-committed-revision');
|
||||
expect(writeOptions.keys[0]).not.toBe(reserveOptions.keys[0]);
|
||||
expect(mockCache.set).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('serializes legacy user catalog migration with Redis-backed writers', async () => {
|
||||
const legacy = { legacy: {} };
|
||||
mockCache.get
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce(legacy)
|
||||
.mockResolvedValueOnce('generation-current')
|
||||
.mockResolvedValueOnce('generation-current');
|
||||
|
||||
await expect(
|
||||
getCachedTools({
|
||||
userId: 'user-1',
|
||||
serverName: 'server-1',
|
||||
configGeneration: 'config-current',
|
||||
}),
|
||||
).resolves.toBe(legacy);
|
||||
|
||||
expect(mockRedisClient.set).toHaveBeenCalledWith(
|
||||
`${CacheKeys.TOOL_CACHE}:tools:mcp-write-lock:user-1:server-1`,
|
||||
expect.any(String),
|
||||
'PX',
|
||||
30_000,
|
||||
'NX',
|
||||
);
|
||||
expect(mockKeyvRedisClient.eval).toHaveBeenCalledWith(
|
||||
expect.stringContaining("redis.call('EXISTS', KEYS[2])"),
|
||||
expect.objectContaining({
|
||||
keys: [
|
||||
`tools:mcp:write-fence:{user-1:server-1}`,
|
||||
`${CacheKeys.TOOL_CACHE}:tools:metadata:mcp:user-legacy-fence:{user-1:server-1}`,
|
||||
`${CacheKeys.TOOL_CACHE}:tools:metadata:mcp:user-generation:{user-1:server-1}`,
|
||||
`${CacheKeys.TOOL_CACHE}:tools:mcp:user:{user-1:server-1}:config-current`,
|
||||
],
|
||||
}),
|
||||
);
|
||||
expect(mockCache.set).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('rejects first-generation creation after its distributed lock ownership is lost', async () => {
|
||||
mockCache.get.mockResolvedValue(null);
|
||||
mockKeyvRedisClient.eval
|
||||
.mockResolvedValueOnce(1)
|
||||
.mockResolvedValueOnce(-1)
|
||||
.mockResolvedValue(1);
|
||||
|
||||
await expect(
|
||||
getMCPToolsCacheGeneration({ userId: 'user-1', serverName: 'server-1' }),
|
||||
).rejects.toThrow('Tool cache lock ownership was lost before generation creation');
|
||||
|
||||
const createCall = mockKeyvRedisClient.eval.mock.calls.find(([script]) =>
|
||||
script.includes("redis.call('EXISTS', KEYS[2])"),
|
||||
);
|
||||
expect(calculateSlot(createCall[1].keys[0])).toBe(calculateSlot(createCall[1].keys[1]));
|
||||
expect(mockCache.set).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('waits through the full abandoned Redis lease before giving up', async () => {
|
||||
jest.useFakeTimers();
|
||||
const startedAt = Date.now();
|
||||
mockRedisClient.set.mockImplementation(async () =>
|
||||
Date.now() - startedAt >= 30_000 ? 'OK' : null,
|
||||
);
|
||||
const operation = jest.fn().mockResolvedValue('recovered');
|
||||
|
||||
const result = runWithGlobalCacheLock(operation);
|
||||
await jest.advanceTimersByTimeAsync(5_000);
|
||||
expect(operation).not.toHaveBeenCalled();
|
||||
|
||||
await jest.advanceTimersByTimeAsync(25_100);
|
||||
await expect(result).resolves.toBe('recovered');
|
||||
expect(operation).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
|
@ -1,4 +1,4 @@
|
|||
const { CacheKeys } = require('librechat-data-provider');
|
||||
const { CacheKeys, Time } = require('librechat-data-provider');
|
||||
|
||||
jest.mock('~/cache/getLogStores');
|
||||
const getLogStores = require('~/cache/getLogStores');
|
||||
|
|
@ -9,76 +9,453 @@ getLogStores.mockReturnValue(mockCache);
|
|||
const {
|
||||
ToolCacheKeys,
|
||||
getCachedTools,
|
||||
updateCachedGlobalTools,
|
||||
setCachedTools,
|
||||
setCachedToolsIfCurrent,
|
||||
getMCPToolsCacheGeneration,
|
||||
renewMCPToolsCacheGeneration,
|
||||
getCachedAppServerTools,
|
||||
getNextAppToolsPublicationRevision,
|
||||
setCachedAppServerTools,
|
||||
runWithGlobalCacheLock,
|
||||
invalidateCachedTools,
|
||||
} = require('../getCachedTools');
|
||||
|
||||
describe('getCachedTools', () => {
|
||||
describe('MCP tool cache', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
getLogStores.mockReturnValue(mockCache);
|
||||
});
|
||||
|
||||
describe('ToolCacheKeys.MCP_SERVER', () => {
|
||||
it('should generate cache keys that include userId', () => {
|
||||
const key = ToolCacheKeys.MCP_SERVER('user123', 'github');
|
||||
expect(key).toBe('tools:mcp:user123:github');
|
||||
it('uses collision-safe configuration-addressed keys', () => {
|
||||
expect(ToolCacheKeys.MCP_APP_SERVER('server:name', 'config/a')).toBe(
|
||||
'tools:mcp:app:server%3Aname:config%2Fa',
|
||||
);
|
||||
expect(ToolCacheKeys.MCP_SERVER('tenant:user', 'server:name', 'config/a')).toBe(
|
||||
'tools:mcp:user:{tenant%3Auser:server%3Aname}:config%2Fa',
|
||||
);
|
||||
expect(ToolCacheKeys.MCP_SERVER('tenant:user', 'server:name', 'config/a')).not.toBe(
|
||||
ToolCacheKeys.MCP_SERVER('tenant', 'user:server:name', 'config/a'),
|
||||
);
|
||||
expect(ToolCacheKeys.MCP_SERVER_GENERATION('tenant:user', 'server:name')).toBe(
|
||||
'tools:metadata:mcp:user-generation:{tenant%3Auser:server%3Aname}',
|
||||
);
|
||||
expect(ToolCacheKeys.MCP_SERVER_GENERATION('tenant:user', 'server:name')).not.toBe(
|
||||
ToolCacheKeys.MCP_SERVER_GENERATION('tenant', 'user:server:name'),
|
||||
);
|
||||
expect(ToolCacheKeys.MCP_SERVER_LEGACY_FENCE('tenant:user', 'server:name')).toBe(
|
||||
'tools:metadata:mcp:user-legacy-fence:{tenant%3Auser:server%3Aname}',
|
||||
);
|
||||
});
|
||||
|
||||
it('keeps the legacy user key available for non-generation callers', () => {
|
||||
expect(ToolCacheKeys.MCP_SERVER('user123', 'github')).toBe('tools:mcp:user123:github');
|
||||
});
|
||||
|
||||
it('gets and sets static global tools without touching MCP slices', async () => {
|
||||
const tools = { builtin: { type: 'function' } };
|
||||
mockCache.get.mockResolvedValue(tools);
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
|
||||
await expect(getCachedTools()).resolves.toBe(tools);
|
||||
await expect(setCachedTools(tools)).resolves.toBe(true);
|
||||
|
||||
expect(mockCache.get).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL);
|
||||
expect(mockCache.set).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL, tools, expect.any(Number));
|
||||
expect(mockCache.delete).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('updates the global catalog atomically through the catalog store', async () => {
|
||||
const current = { builtin: { type: 'function' }, old_mcp_server: { type: 'function' } };
|
||||
mockCache.get.mockResolvedValue(current);
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
|
||||
await updateCachedGlobalTools(({ old_mcp_server: _removed, ...staticTools }) => staticTools);
|
||||
|
||||
expect(mockCache.get).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL);
|
||||
expect(mockCache.set).toHaveBeenCalledWith(
|
||||
ToolCacheKeys.GLOBAL,
|
||||
{ builtin: { type: 'function' } },
|
||||
Time.TWELVE_HOURS,
|
||||
);
|
||||
});
|
||||
|
||||
it('recreates the authoritative global catalog after its cache entry expires', async () => {
|
||||
mockCache.get.mockResolvedValue(null);
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
|
||||
await updateCachedGlobalTools(() => ({ builtin: { type: 'function' } }));
|
||||
|
||||
expect(mockCache.set).toHaveBeenCalledWith(
|
||||
ToolCacheKeys.GLOBAL,
|
||||
{ builtin: { type: 'function' } },
|
||||
Time.TWELVE_HOURS,
|
||||
);
|
||||
});
|
||||
|
||||
it('gets and sets an authoritative app slice, including an empty catalog', async () => {
|
||||
mockCache.get.mockResolvedValue({});
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
|
||||
await expect(getCachedAppServerTools('github', 'config-v2')).resolves.toEqual({});
|
||||
await expect(setCachedAppServerTools('github', 'config-v2', {})).resolves.toBe(true);
|
||||
|
||||
const key = ToolCacheKeys.MCP_APP_SERVER('github', 'config-v2');
|
||||
expect(mockCache.get).toHaveBeenCalledWith(key);
|
||||
expect(mockCache.set).toHaveBeenCalledWith(
|
||||
key,
|
||||
{ version: 1, publicationRevision: '0', tools: {} },
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
|
||||
it('prevents a slow older app snapshot from replacing a newer revision', async () => {
|
||||
const key = ToolCacheKeys.MCP_APP_SERVER('github', 'config-v2');
|
||||
let cached = null;
|
||||
mockCache.get.mockImplementation(async (requestedKey) =>
|
||||
requestedKey === key ? cached : null,
|
||||
);
|
||||
mockCache.set.mockImplementation(async (requestedKey, value) => {
|
||||
if (requestedKey === key) {
|
||||
cached = value;
|
||||
}
|
||||
return true;
|
||||
});
|
||||
|
||||
const older = await getNextAppToolsPublicationRevision('github', 'config-v2');
|
||||
const newer = await getNextAppToolsPublicationRevision('github', 'config-v2');
|
||||
const currentTools = { current: { type: 'function' } };
|
||||
const staleTools = { stale: { type: 'function' } };
|
||||
|
||||
await expect(setCachedAppServerTools('github', 'config-v2', currentTools, newer)).resolves.toBe(
|
||||
true,
|
||||
);
|
||||
await expect(setCachedAppServerTools('github', 'config-v2', staleTools, older)).resolves.toBe(
|
||||
false,
|
||||
);
|
||||
await expect(getCachedAppServerTools('github', 'config-v2')).resolves.toEqual(currentTools);
|
||||
});
|
||||
|
||||
it('allows a completed snapshot when a later reserved request aborts', async () => {
|
||||
const key = ToolCacheKeys.MCP_APP_SERVER('github', 'config-v2');
|
||||
let cached = null;
|
||||
mockCache.get.mockImplementation(async (requestedKey) =>
|
||||
requestedKey === key ? cached : null,
|
||||
);
|
||||
mockCache.set.mockImplementation(async (requestedKey, value) => {
|
||||
if (requestedKey === key) {
|
||||
cached = value;
|
||||
}
|
||||
return true;
|
||||
});
|
||||
|
||||
const completed = await getNextAppToolsPublicationRevision('github', 'config-v2');
|
||||
await getNextAppToolsPublicationRevision('github', 'config-v2');
|
||||
|
||||
await expect(
|
||||
setCachedAppServerTools('github', 'config-v2', { completed: {} }, completed),
|
||||
).resolves.toBe(true);
|
||||
await expect(getCachedAppServerTools('github', 'config-v2')).resolves.toEqual({
|
||||
completed: {},
|
||||
});
|
||||
});
|
||||
|
||||
describe('TOOL_CACHE namespace usage', () => {
|
||||
it('getCachedTools should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.get.mockResolvedValue(null);
|
||||
await getCachedTools();
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
it('stores unguarded user tools under the supplied config generation', async () => {
|
||||
const tools = { search: { type: 'function' } };
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
|
||||
await setCachedTools(tools, {
|
||||
userId: 'user1',
|
||||
serverName: 'github',
|
||||
configGeneration: 'config-v2',
|
||||
});
|
||||
|
||||
it('getCachedTools with MCP server options should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.get.mockResolvedValue({ tool1: {} });
|
||||
await getCachedTools({ userId: 'user1', serverName: 'github' });
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
expect(mockCache.get).toHaveBeenCalledWith(ToolCacheKeys.MCP_SERVER('user1', 'github'));
|
||||
});
|
||||
expect(mockCache.set).toHaveBeenCalledWith(
|
||||
ToolCacheKeys.MCP_SERVER('user1', 'github', 'config-v2'),
|
||||
tools,
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
|
||||
it('setCachedTools should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
const tools = { tool1: { type: 'function' } };
|
||||
await setCachedTools(tools);
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
expect(mockCache.set).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL, tools, expect.any(Number));
|
||||
});
|
||||
it('writes guarded user tools under both config and connection generations', async () => {
|
||||
const tools = { current: { type: 'function' } };
|
||||
mockCache.get.mockResolvedValue('connection-a');
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
|
||||
it('setCachedTools with MCP server options should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
const tools = { tool1: { type: 'function' } };
|
||||
await setCachedTools(tools, { userId: 'user1', serverName: 'github' });
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
expect(mockCache.set).toHaveBeenCalledWith(
|
||||
ToolCacheKeys.MCP_SERVER('user1', 'github'),
|
||||
await expect(
|
||||
setCachedToolsIfCurrent(tools, {
|
||||
userId: 'user1',
|
||||
serverName: 'github',
|
||||
configGeneration: 'config-v2',
|
||||
publicationGeneration: 'connection-a',
|
||||
}),
|
||||
).resolves.toBe(true);
|
||||
|
||||
expect(mockCache.set).toHaveBeenCalledWith(
|
||||
ToolCacheKeys.MCP_SERVER('user1', 'github', 'config-v2'),
|
||||
{ version: 1, publicationGeneration: 'connection-a', tools },
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
|
||||
it('reads guarded user tools only while their connection generation is current', async () => {
|
||||
const tools = { current: { type: 'function' } };
|
||||
mockCache.get
|
||||
.mockResolvedValueOnce({
|
||||
version: 1,
|
||||
publicationGeneration: 'connection-a',
|
||||
tools,
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
})
|
||||
.mockResolvedValueOnce('connection-a');
|
||||
|
||||
it('invalidateCachedTools should use TOOL_CACHE namespace', async () => {
|
||||
mockCache.delete.mockResolvedValue(true);
|
||||
await invalidateCachedTools({ invalidateGlobal: true });
|
||||
expect(getLogStores).toHaveBeenCalledWith(CacheKeys.TOOL_CACHE);
|
||||
expect(mockCache.delete).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL);
|
||||
});
|
||||
await expect(
|
||||
getCachedTools({
|
||||
userId: 'user1',
|
||||
serverName: 'github',
|
||||
configGeneration: 'config-v2',
|
||||
}),
|
||||
).resolves.toEqual(tools);
|
||||
|
||||
it('should NOT use CONFIG_STORE namespace', async () => {
|
||||
mockCache.get.mockResolvedValue(null);
|
||||
await getCachedTools();
|
||||
await getCachedTools({ userId: 'user1', serverName: 'github' });
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
await setCachedTools({ tool1: {} });
|
||||
mockCache.delete.mockResolvedValue(true);
|
||||
await invalidateCachedTools({ invalidateGlobal: true });
|
||||
expect(mockCache.get.mock.calls[0][0]).toBe(
|
||||
ToolCacheKeys.MCP_SERVER('user1', 'github', 'config-v2'),
|
||||
);
|
||||
});
|
||||
|
||||
const allCalls = getLogStores.mock.calls.flat();
|
||||
expect(allCalls).not.toContain(CacheKeys.CONFIG_STORE);
|
||||
expect(allCalls.every((key) => key === CacheKeys.TOOL_CACHE)).toBe(true);
|
||||
});
|
||||
it('copies a legacy user catalog into the config-addressed key on rollout', async () => {
|
||||
const tools = { legacy: { type: 'function' } };
|
||||
mockCache.get
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce(tools)
|
||||
.mockResolvedValueOnce('connection-a')
|
||||
.mockResolvedValueOnce('connection-a');
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
|
||||
await expect(
|
||||
getCachedTools({
|
||||
userId: 'user1',
|
||||
serverName: 'github',
|
||||
configGeneration: 'config-v2',
|
||||
}),
|
||||
).resolves.toBe(tools);
|
||||
|
||||
expect(mockCache.get).toHaveBeenNthCalledWith(4, ToolCacheKeys.MCP_SERVER('user1', 'github'));
|
||||
expect(mockCache.set).toHaveBeenCalledWith(
|
||||
ToolCacheKeys.MCP_SERVER('user1', 'github', 'config-v2'),
|
||||
{
|
||||
version: 1,
|
||||
publicationGeneration: 'connection-a',
|
||||
tools,
|
||||
},
|
||||
Time.TWELVE_HOURS,
|
||||
);
|
||||
});
|
||||
|
||||
it('creates a generation fence before migrating a legacy user catalog', async () => {
|
||||
const tools = { legacy: { type: 'function' } };
|
||||
mockCache.get
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce(tools)
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockImplementationOnce(async () => mockCache.set.mock.calls[0][1]);
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
|
||||
await expect(
|
||||
getCachedTools({
|
||||
userId: 'user1',
|
||||
serverName: 'github',
|
||||
configGeneration: 'config-v2',
|
||||
}),
|
||||
).resolves.toBe(tools);
|
||||
|
||||
const [generationKey, generation] = mockCache.set.mock.calls[0];
|
||||
expect(generationKey).toBe(ToolCacheKeys.MCP_SERVER_GENERATION('user1', 'github'));
|
||||
expect(generation).toEqual(expect.any(String));
|
||||
expect(mockCache.set.mock.calls[1][1]).toEqual(
|
||||
expect.objectContaining({ publicationGeneration: generation, tools }),
|
||||
);
|
||||
});
|
||||
|
||||
it('rejects a legacy catalog recreated after the scope was fenced', async () => {
|
||||
mockCache.get
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce(true);
|
||||
|
||||
await expect(
|
||||
getCachedTools({
|
||||
userId: 'user1',
|
||||
serverName: 'github',
|
||||
configGeneration: 'config-v2',
|
||||
}),
|
||||
).resolves.toBeNull();
|
||||
|
||||
expect(mockCache.get).not.toHaveBeenCalledWith(ToolCacheKeys.MCP_SERVER('user1', 'github'));
|
||||
expect(mockCache.set).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('preserves a config-addressed catalog published during legacy fallback', async () => {
|
||||
const current = { current: { type: 'function' } };
|
||||
mockCache.get.mockResolvedValueOnce(null).mockResolvedValueOnce(current);
|
||||
|
||||
await expect(
|
||||
getCachedTools({
|
||||
userId: 'user1',
|
||||
serverName: 'github',
|
||||
configGeneration: 'config-v2',
|
||||
}),
|
||||
).resolves.toBe(current);
|
||||
|
||||
expect(mockCache.get).toHaveBeenCalledTimes(2);
|
||||
expect(mockCache.set).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('hides a guarded entry after its connection generation is replaced', async () => {
|
||||
mockCache.get
|
||||
.mockResolvedValueOnce({
|
||||
version: 1,
|
||||
publicationGeneration: 'connection-a',
|
||||
tools: { stale: {} },
|
||||
})
|
||||
.mockResolvedValueOnce('connection-b');
|
||||
|
||||
await expect(
|
||||
getCachedTools({
|
||||
userId: 'user1',
|
||||
serverName: 'github',
|
||||
configGeneration: 'config-v1',
|
||||
}),
|
||||
).resolves.toBeNull();
|
||||
});
|
||||
|
||||
it('cannot let a late old-config write replace the current config key', async () => {
|
||||
mockCache.get.mockResolvedValue('connection-a');
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
|
||||
await setCachedToolsIfCurrent(
|
||||
{ current: {} },
|
||||
{
|
||||
userId: 'user1',
|
||||
serverName: 'github',
|
||||
configGeneration: 'config-v2',
|
||||
publicationGeneration: 'connection-a',
|
||||
},
|
||||
);
|
||||
await setCachedToolsIfCurrent(
|
||||
{ stale: {} },
|
||||
{
|
||||
userId: 'user1',
|
||||
serverName: 'github',
|
||||
configGeneration: 'config-v1',
|
||||
publicationGeneration: 'connection-a',
|
||||
},
|
||||
);
|
||||
|
||||
expect(mockCache.set).toHaveBeenCalledWith(
|
||||
ToolCacheKeys.MCP_SERVER('user1', 'github', 'config-v2'),
|
||||
expect.objectContaining({ tools: { current: {} } }),
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(mockCache.set).toHaveBeenCalledWith(
|
||||
ToolCacheKeys.MCP_SERVER('user1', 'github', 'config-v1'),
|
||||
expect.objectContaining({ tools: { stale: {} } }),
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
|
||||
it('creates and reuses a durable connection publication generation', async () => {
|
||||
mockCache.get
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce(null)
|
||||
.mockResolvedValueOnce('existing-generation');
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
|
||||
const created = await getMCPToolsCacheGeneration({ userId: 'user1', serverName: 'github' });
|
||||
const existing = await getMCPToolsCacheGeneration({ userId: 'user1', serverName: 'github' });
|
||||
|
||||
expect(created).toEqual(expect.any(String));
|
||||
expect(existing).toBe('existing-generation');
|
||||
expect(mockCache.set).toHaveBeenCalledWith(
|
||||
ToolCacheKeys.MCP_SERVER_GENERATION('user1', 'github'),
|
||||
created,
|
||||
expect.any(Number),
|
||||
);
|
||||
});
|
||||
|
||||
it('renews a lease only for its current publication generation', async () => {
|
||||
mockCache.get.mockResolvedValue('connection-a');
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
|
||||
await expect(
|
||||
renewMCPToolsCacheGeneration({
|
||||
userId: 'user1',
|
||||
serverName: 'github',
|
||||
publicationGeneration: 'connection-a',
|
||||
}),
|
||||
).resolves.toBe(true);
|
||||
await expect(
|
||||
renewMCPToolsCacheGeneration({
|
||||
userId: 'user1',
|
||||
serverName: 'github',
|
||||
publicationGeneration: 'connection-b',
|
||||
}),
|
||||
).resolves.toBe(false);
|
||||
});
|
||||
|
||||
it('rotates the connection generation before deleting the legacy user key', async () => {
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
mockCache.delete.mockResolvedValue(true);
|
||||
|
||||
await invalidateCachedTools({ userId: 'user1', serverName: 'github' });
|
||||
|
||||
expect(mockCache.set).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
ToolCacheKeys.MCP_SERVER_LEGACY_FENCE('user1', 'github'),
|
||||
true,
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(mockCache.set.mock.calls[0][2]).toBeGreaterThanOrEqual(Time.ONE_DAY);
|
||||
expect(mockCache.set).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
ToolCacheKeys.MCP_SERVER_GENERATION('user1', 'github'),
|
||||
expect.any(String),
|
||||
expect.any(Number),
|
||||
);
|
||||
expect(mockCache.set.mock.calls[1][2]).toBeGreaterThanOrEqual(Time.ONE_DAY);
|
||||
expect(mockCache.delete).toHaveBeenCalledWith(ToolCacheKeys.MCP_SERVER('user1', 'github'));
|
||||
expect(mockCache.set.mock.invocationCallOrder[1]).toBeLessThan(
|
||||
mockCache.delete.mock.invocationCallOrder[0],
|
||||
);
|
||||
});
|
||||
|
||||
it('invalidates only the static global key for broad config changes', async () => {
|
||||
mockCache.delete.mockResolvedValue(true);
|
||||
|
||||
await invalidateCachedTools({ invalidateGlobal: true });
|
||||
|
||||
expect(mockCache.delete).toHaveBeenCalledTimes(1);
|
||||
expect(mockCache.delete).toHaveBeenCalledWith(ToolCacheKeys.GLOBAL);
|
||||
});
|
||||
|
||||
it('runs global cache operations directly when the cache is in memory', async () => {
|
||||
const operation = jest.fn().mockResolvedValue('done');
|
||||
await expect(runWithGlobalCacheLock(operation)).resolves.toBe('done');
|
||||
expect(operation).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('uses only the TOOL_CACHE namespace', async () => {
|
||||
mockCache.get.mockResolvedValue(null);
|
||||
mockCache.set.mockResolvedValue(true);
|
||||
mockCache.delete.mockResolvedValue(true);
|
||||
|
||||
await getCachedTools();
|
||||
await getCachedAppServerTools('github', 'config-v2');
|
||||
await setCachedTools({});
|
||||
await invalidateCachedTools({ invalidateGlobal: true });
|
||||
|
||||
expect(getLogStores.mock.calls.flat().every((key) => key === CacheKeys.TOOL_CACHE)).toBe(true);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,89 +1,20 @@
|
|||
const { CacheKeys, Time } = require('librechat-data-provider');
|
||||
const { CacheKeys } = require('librechat-data-provider');
|
||||
const {
|
||||
cacheConfig,
|
||||
ioredisClient,
|
||||
keyvRedisClient,
|
||||
mcpConfig,
|
||||
ToolCacheKeys,
|
||||
createMCPCatalogStore,
|
||||
} = require('@librechat/api');
|
||||
const getLogStores = require('~/cache/getLogStores');
|
||||
|
||||
/**
|
||||
* Cache key generators for different tool access patterns
|
||||
*/
|
||||
const ToolCacheKeys = {
|
||||
/** Global tools available to all users */
|
||||
GLOBAL: 'tools:global',
|
||||
/** MCP tools cached by user ID and server name */
|
||||
MCP_SERVER: (userId, serverName) => `tools:mcp:${userId}:${serverName}`,
|
||||
};
|
||||
const store = createMCPCatalogStore({
|
||||
cacheConfig,
|
||||
ioredisClient,
|
||||
keyvRedisClient,
|
||||
userConnectionIdleTimeout: mcpConfig.USER_CONNECTION_IDLE_TIMEOUT,
|
||||
getCache: () => getLogStores(CacheKeys.TOOL_CACHE),
|
||||
});
|
||||
|
||||
/**
|
||||
* Retrieves available tools from cache
|
||||
* @function getCachedTools
|
||||
* @param {Object} options - Options for retrieving tools
|
||||
* @param {string} [options.userId] - User ID for user-specific MCP tools
|
||||
* @param {string} [options.serverName] - MCP server name to get cached tools for
|
||||
* @returns {Promise<LCAvailableTools|null>} The available tools object or null if not cached
|
||||
*/
|
||||
async function getCachedTools(options = {}) {
|
||||
const cache = getLogStores(CacheKeys.TOOL_CACHE);
|
||||
const { userId, serverName } = options;
|
||||
|
||||
// Return MCP server-specific tools if requested
|
||||
if (serverName && userId) {
|
||||
return await cache.get(ToolCacheKeys.MCP_SERVER(userId, serverName));
|
||||
}
|
||||
|
||||
// Default to global tools
|
||||
return await cache.get(ToolCacheKeys.GLOBAL);
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets available tools in cache
|
||||
* @function setCachedTools
|
||||
* @param {Object} tools - The tools object to cache
|
||||
* @param {Object} options - Options for caching tools
|
||||
* @param {string} [options.userId] - User ID for user-specific MCP tools
|
||||
* @param {string} [options.serverName] - MCP server name for server-specific tools
|
||||
* @param {number} [options.ttl] - Time to live in milliseconds (default: 12 hours)
|
||||
* @returns {Promise<boolean>} Whether the operation was successful
|
||||
*/
|
||||
async function setCachedTools(tools, options = {}) {
|
||||
const cache = getLogStores(CacheKeys.TOOL_CACHE);
|
||||
const { userId, serverName, ttl = Time.TWELVE_HOURS } = options;
|
||||
|
||||
// Cache by MCP server if specified (requires userId)
|
||||
if (serverName && userId) {
|
||||
return await cache.set(ToolCacheKeys.MCP_SERVER(userId, serverName), tools, ttl);
|
||||
}
|
||||
|
||||
// Default to global cache
|
||||
return await cache.set(ToolCacheKeys.GLOBAL, tools, ttl);
|
||||
}
|
||||
|
||||
/**
|
||||
* Invalidates cached tools
|
||||
* @function invalidateCachedTools
|
||||
* @param {Object} options - Options for invalidating tools
|
||||
* @param {string} [options.userId] - User ID for user-specific MCP tools
|
||||
* @param {string} [options.serverName] - MCP server name to invalidate
|
||||
* @param {boolean} [options.invalidateGlobal=false] - Whether to invalidate global tools
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
async function invalidateCachedTools(options = {}) {
|
||||
const cache = getLogStores(CacheKeys.TOOL_CACHE);
|
||||
const { userId, serverName, invalidateGlobal = false } = options;
|
||||
|
||||
const keysToDelete = [];
|
||||
|
||||
if (invalidateGlobal) {
|
||||
keysToDelete.push(ToolCacheKeys.GLOBAL);
|
||||
}
|
||||
|
||||
if (serverName && userId) {
|
||||
keysToDelete.push(ToolCacheKeys.MCP_SERVER(userId, serverName));
|
||||
}
|
||||
|
||||
await Promise.all(keysToDelete.map((key) => cache.delete(key)));
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
ToolCacheKeys,
|
||||
getCachedTools,
|
||||
setCachedTools,
|
||||
invalidateCachedTools,
|
||||
};
|
||||
module.exports = { ToolCacheKeys, ...store };
|
||||
|
|
|
|||
|
|
@ -1,17 +1,43 @@
|
|||
const { createMCPToolCacheService, MCPServersRegistry } = require('@librechat/api');
|
||||
const { getCachedTools, setCachedTools } = require('./getCachedTools');
|
||||
const {
|
||||
getCachedTools,
|
||||
updateCachedGlobalTools,
|
||||
setCachedToolsWithinGlobalLock,
|
||||
getCachedAppServerTools,
|
||||
setCachedAppServerTools,
|
||||
setCachedToolsIfCurrent,
|
||||
getMCPToolsCacheGeneration,
|
||||
renewMCPToolsCacheGeneration,
|
||||
getNextAppToolsPublicationRevision,
|
||||
} = require('./getCachedTools');
|
||||
|
||||
const { mergeAppTools, cacheMCPServerTools, updateMCPServerTools, getMCPServerTools } =
|
||||
createMCPToolCacheService({
|
||||
getCachedTools,
|
||||
setCachedTools,
|
||||
getServerConfig: (serverName, userId) =>
|
||||
MCPServersRegistry.getInstance().getServerConfig(serverName, userId),
|
||||
});
|
||||
const {
|
||||
syncStaticTools,
|
||||
mergeAppTools,
|
||||
cacheMCPServerTools,
|
||||
updateMCPServerTools,
|
||||
getMCPServerTools,
|
||||
} = createMCPToolCacheService({
|
||||
getCachedTools,
|
||||
updateCachedGlobalTools,
|
||||
setCachedTools: setCachedToolsWithinGlobalLock,
|
||||
setCachedToolsIfCurrent,
|
||||
getCachedAppServerTools,
|
||||
setCachedAppServerTools,
|
||||
getServerConfig: (serverName, userId) =>
|
||||
MCPServersRegistry.getInstance().getServerConfig(serverName, userId),
|
||||
getAllServerConfigs: () => MCPServersRegistry.getInstance().getAllServerConfigs(),
|
||||
isAppServerConfig: (serverName, effectiveConfig) =>
|
||||
MCPServersRegistry.getInstance().isAppServerConfig(serverName, effectiveConfig),
|
||||
});
|
||||
|
||||
module.exports = {
|
||||
syncStaticTools,
|
||||
mergeAppTools,
|
||||
getMCPServerTools,
|
||||
cacheMCPServerTools,
|
||||
updateMCPServerTools,
|
||||
getMCPToolsCacheGeneration,
|
||||
renewMCPToolsCacheGeneration,
|
||||
getNextAppToolsPublicationRevision,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ const {
|
|||
normalizeMCPToolKey,
|
||||
buildServerNameAliases,
|
||||
findShadowedServerNames,
|
||||
getAssistantToolDefinitions: loadAssistantToolDefinitions,
|
||||
resolveMCPServerContext,
|
||||
normalizeJsonSchema,
|
||||
GenerationJobManager,
|
||||
|
|
@ -27,6 +28,7 @@ const {
|
|||
isUserSourced,
|
||||
checkAccessWithRequestCache,
|
||||
getMissingCustomUserVars,
|
||||
getUserMCPAuthMap,
|
||||
getServerCustomUserVars,
|
||||
requiresEphemeralUserConnection,
|
||||
requiresOAuthMachinery,
|
||||
|
|
@ -49,12 +51,17 @@ const {
|
|||
getMCPManager,
|
||||
} = require('~/config');
|
||||
const db = require('~/models');
|
||||
const { findToken, createToken, updateToken, deleteTokens } = db;
|
||||
const { findToken, createToken, updateToken, deleteTokens, findPluginAuthsByKeys } = db;
|
||||
const { getGraphApiToken } = require('./GraphTokenService');
|
||||
const { exchangeOboToken } = require('./OboTokenService');
|
||||
const { createOboTrustChecker } = require('./OboPolicyService');
|
||||
const { reinitMCPServer } = require('./Tools/mcp');
|
||||
const { getAppConfig } = require('./Config');
|
||||
const {
|
||||
getAppConfig,
|
||||
getCachedTools,
|
||||
getMCPServerTools,
|
||||
cacheMCPServerTools,
|
||||
} = require('./Config');
|
||||
const { getLogStores } = require('~/cache');
|
||||
|
||||
const MAX_CACHE_SIZE = 1000;
|
||||
|
|
@ -303,6 +310,57 @@ async function healMcpToolNames({ req, tools, toolDefinitions }) {
|
|||
return healedList;
|
||||
}
|
||||
|
||||
/**
|
||||
* Loads static and MCP function definitions used by assistant create/update writes. MCP catalogs
|
||||
* are stored per server and effective config, so assistant writers must resolve the referenced
|
||||
* server slices instead of relying on the static aggregate cache.
|
||||
* @param {object} params
|
||||
* @param {ServerRequest} params.req
|
||||
* @param {Array<string | object>} [params.tools]
|
||||
* @returns {Promise<object>}
|
||||
*/
|
||||
async function getAssistantToolDefinitions({ req, tools }) {
|
||||
const registry = getMCPServersRegistry();
|
||||
const appConfig = await getAppConfigForRequest(req);
|
||||
return await loadAssistantToolDefinitions(
|
||||
{
|
||||
user: req.user,
|
||||
tools,
|
||||
staticTools: (await getCachedTools()) ?? {},
|
||||
mcpConfig: appConfig?.mcpConfig ?? {},
|
||||
},
|
||||
{
|
||||
ensureConfigServers: (mcpConfig) => registry.ensureConfigServers(mcpConfig),
|
||||
getAllServerConfigs: (userId, configServers, role) =>
|
||||
registry.getAllServerConfigs(userId, configServers, role),
|
||||
getMCPServerTools,
|
||||
getServerToolFunctionsSnapshot: async (userId, serverName, serverConfig) =>
|
||||
(await getMCPManager()?.getServerToolFunctionsSnapshot(
|
||||
userId,
|
||||
serverName,
|
||||
serverConfig,
|
||||
)) ?? {
|
||||
tools: null,
|
||||
},
|
||||
recoverServerTools: async (serverName, serverConfig) => {
|
||||
const userMCPAuthMap = await getUserMCPAuthMap({
|
||||
userId: req.user.id,
|
||||
servers: [serverName],
|
||||
findPluginAuthsByKeys,
|
||||
});
|
||||
const result = await reinitMCPServer({
|
||||
user: req.user,
|
||||
serverName,
|
||||
serverConfig,
|
||||
userMCPAuthMap,
|
||||
});
|
||||
return result?.availableTools ?? null;
|
||||
},
|
||||
cacheMCPServerTools,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves the name set MCP collision guards audit against. Prefers the
|
||||
* caller-threaded accessible set; self-fetches only when a configured name
|
||||
|
|
@ -1431,6 +1489,7 @@ module.exports = {
|
|||
resolveMcpServerContext,
|
||||
getAccessibleMcpServerNames,
|
||||
healMcpToolNames,
|
||||
getAssistantToolDefinitions,
|
||||
resolveCollisionAuditNames,
|
||||
resolveMcpConfigNames,
|
||||
resolveAllMcpConfigs,
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ 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 { getMCPToolsCacheGeneration, updateMCPServerTools } = require('~/server/services/Config');
|
||||
const { getLogStores } = require('~/cache');
|
||||
|
||||
const MCP_REINITIALIZE_FAILURE_REASONS = {
|
||||
|
|
@ -65,6 +65,7 @@ async function reinitMCPServer({
|
|||
let oauthUrl = null;
|
||||
let oauthExpiresAt;
|
||||
let ephemeralServer = false;
|
||||
let publicationGeneration;
|
||||
|
||||
try {
|
||||
const registry = getMCPServersRegistry();
|
||||
|
|
@ -167,6 +168,13 @@ async function reinitMCPServer({
|
|||
const mcpManager = getMCPManager();
|
||||
const tokenMethods = { findToken, updateToken, createToken, deleteTokens };
|
||||
|
||||
if (!ephemeralServer) {
|
||||
publicationGeneration = await getMCPToolsCacheGeneration({
|
||||
userId: user.id,
|
||||
serverName,
|
||||
});
|
||||
}
|
||||
|
||||
const oauthStart =
|
||||
_oauthStart ??
|
||||
(async (authURL, options) => {
|
||||
|
|
@ -259,7 +267,36 @@ async function reinitMCPServer({
|
|||
}
|
||||
|
||||
if (connection && !oauthRequired) {
|
||||
tools = await connection.fetchTools();
|
||||
publicationGeneration =
|
||||
mcpManager.getToolPublicationGeneration(connection) ?? publicationGeneration;
|
||||
let snapshot;
|
||||
if (typeof connection.fetchOrderedToolsSnapshot === 'function') {
|
||||
snapshot = await connection.fetchOrderedToolsSnapshot();
|
||||
} else if (typeof connection.fetchToolsSnapshot === 'function') {
|
||||
snapshot = await connection.fetchToolsSnapshot();
|
||||
} else {
|
||||
snapshot = { tools: await connection.fetchTools(), complete: true };
|
||||
}
|
||||
if (snapshot.complete) {
|
||||
tools = snapshot.tools;
|
||||
} else {
|
||||
logger.warn(
|
||||
`[MCP Reinitialize] Preserving cached tools for ${serverName} because tools/list returned an incomplete snapshot`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
if (tools && !ephemeralServer && publicationGeneration) {
|
||||
const currentGeneration = await getMCPToolsCacheGeneration({
|
||||
userId: user.id,
|
||||
serverName,
|
||||
});
|
||||
if (currentGeneration !== publicationGeneration) {
|
||||
logger.warn(
|
||||
`[MCP Reinitialize] Discarding stale tools for ${serverName} because its publication generation changed during discovery`,
|
||||
);
|
||||
tools = null;
|
||||
}
|
||||
}
|
||||
|
||||
if (tools) {
|
||||
|
|
@ -268,7 +305,11 @@ async function reinitMCPServer({
|
|||
serverName,
|
||||
tools,
|
||||
serverConfig,
|
||||
...(publicationGeneration && { publicationGeneration }),
|
||||
});
|
||||
if (availableTools == null) {
|
||||
tools = null;
|
||||
}
|
||||
}
|
||||
|
||||
logger.debug(
|
||||
|
|
@ -325,12 +366,9 @@ async function reinitMCPServer({
|
|||
} finally {
|
||||
if (connection && ephemeralServer && !requestScopedConnections) {
|
||||
try {
|
||||
await connection.disconnect();
|
||||
await connection.dispose();
|
||||
} catch (error) {
|
||||
logger.warn(
|
||||
`[MCP Reinitialize] Failed to disconnect ephemeral server ${serverName}`,
|
||||
error,
|
||||
);
|
||||
logger.warn(`[MCP Reinitialize] Failed to dispose ephemeral server ${serverName}`, error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,11 +4,14 @@ const mockGetConnection = jest.fn();
|
|||
const mockDiscoverServerTools = jest.fn();
|
||||
const mockGetGraphApiToken = jest.fn();
|
||||
const mockUpdateMCPServerTools = jest.fn();
|
||||
const mockGetMCPToolsCacheGeneration = jest.fn().mockResolvedValue('generation-current');
|
||||
const mockGetToolPublicationGeneration = jest.fn().mockReturnValue('generation-current');
|
||||
|
||||
jest.mock('~/config', () => ({
|
||||
getMCPManager: jest.fn(() => ({
|
||||
getConnection: mockGetConnection,
|
||||
discoverServerTools: mockDiscoverServerTools,
|
||||
getToolPublicationGeneration: mockGetToolPublicationGeneration,
|
||||
})),
|
||||
getMCPServersRegistry: jest.fn(() => ({ getServerConfig: jest.fn() })),
|
||||
getFlowStateManager: jest.fn(() => ({})),
|
||||
|
|
@ -21,6 +24,7 @@ jest.mock('~/models', () => ({
|
|||
}));
|
||||
jest.mock('~/server/services/Config', () => ({
|
||||
updateMCPServerTools: mockUpdateMCPServerTools,
|
||||
getMCPToolsCacheGeneration: mockGetMCPToolsCacheGeneration,
|
||||
}));
|
||||
jest.mock('~/server/services/GraphTokenService', () => ({
|
||||
getGraphApiToken: mockGetGraphApiToken,
|
||||
|
|
@ -117,9 +121,71 @@ describe('reinitMCPServer — customUserVars gating (issue #10969)', () => {
|
|||
serverName,
|
||||
tools: [],
|
||||
serverConfig: { type: 'streamable-http', url: 'https://thingy.example.com/mcp' },
|
||||
publicationGeneration: 'generation-current',
|
||||
});
|
||||
});
|
||||
|
||||
it('preserves cached tools when live recovery returns an incomplete snapshot', async () => {
|
||||
const fetchOrderedToolsSnapshot = jest.fn().mockResolvedValue({
|
||||
tools: [{ name: 'partial', inputSchema: { type: 'object' } }],
|
||||
complete: false,
|
||||
});
|
||||
mockGetConnection.mockResolvedValue({
|
||||
fetchOrderedToolsSnapshot,
|
||||
});
|
||||
|
||||
const result = await reinitMCPServer({
|
||||
user,
|
||||
serverName,
|
||||
serverConfig: { type: 'streamable-http', url: 'https://thingy.example.com/mcp' },
|
||||
});
|
||||
|
||||
expect(result.tools).toBeNull();
|
||||
expect(fetchOrderedToolsSnapshot).toHaveBeenCalledTimes(1);
|
||||
expect(mockUpdateMCPServerTools).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('discards a snapshot when another replica rotates its generation during discovery', async () => {
|
||||
mockGetMCPToolsCacheGeneration
|
||||
.mockResolvedValueOnce('generation-current')
|
||||
.mockResolvedValueOnce('generation-replaced');
|
||||
mockGetConnection.mockResolvedValue({
|
||||
fetchOrderedToolsSnapshot: jest.fn().mockResolvedValue({
|
||||
tools: [{ name: 'stale', inputSchema: { type: 'object' } }],
|
||||
complete: true,
|
||||
}),
|
||||
});
|
||||
|
||||
const result = await reinitMCPServer({
|
||||
user,
|
||||
serverName,
|
||||
serverConfig: { type: 'streamable-http', url: 'https://thingy.example.com/mcp' },
|
||||
});
|
||||
|
||||
expect(result.tools).toBeNull();
|
||||
expect(result.availableTools).toBeNull();
|
||||
expect(mockUpdateMCPServerTools).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('does not return tools when the guarded publication loses its generation race', async () => {
|
||||
mockGetConnection.mockResolvedValue({
|
||||
fetchOrderedToolsSnapshot: jest.fn().mockResolvedValue({
|
||||
tools: [{ name: 'stale', inputSchema: { type: 'object' } }],
|
||||
complete: true,
|
||||
}),
|
||||
});
|
||||
mockUpdateMCPServerTools.mockResolvedValue(null);
|
||||
|
||||
const result = await reinitMCPServer({
|
||||
user,
|
||||
serverName,
|
||||
serverConfig: { type: 'streamable-http', url: 'https://thingy.example.com/mcp' },
|
||||
});
|
||||
|
||||
expect(result.tools).toBeNull();
|
||||
expect(result.availableTools).toBeNull();
|
||||
});
|
||||
|
||||
it('passes request body and Graph resolver into connection creation', async () => {
|
||||
mockGetConnection.mockResolvedValue({ fetchTools: jest.fn().mockResolvedValue([]) });
|
||||
const requestBody = { conversationId: 'conv-123', messageId: 'msg-123' };
|
||||
|
|
@ -167,8 +233,8 @@ describe('reinitMCPServer — customUserVars gating (issue #10969)', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('disconnects ephemeral BODY-scoped connections after loading tools', async () => {
|
||||
const disconnect = jest.fn().mockResolvedValue(undefined);
|
||||
it('disposes ephemeral BODY-scoped connections after loading tools', async () => {
|
||||
const dispose = jest.fn().mockResolvedValue(undefined);
|
||||
const tools = [{ name: 'search', inputSchema: { type: 'object', properties: {} } }];
|
||||
const serverConfig = {
|
||||
type: 'streamable-http',
|
||||
|
|
@ -176,7 +242,7 @@ describe('reinitMCPServer — customUserVars gating (issue #10969)', () => {
|
|||
source: 'yaml',
|
||||
};
|
||||
mockGetConnection.mockResolvedValue({
|
||||
disconnect,
|
||||
dispose,
|
||||
fetchTools: jest.fn().mockResolvedValue(tools),
|
||||
});
|
||||
|
||||
|
|
@ -188,7 +254,7 @@ describe('reinitMCPServer — customUserVars gating (issue #10969)', () => {
|
|||
userMCPAuthMap: undefined,
|
||||
});
|
||||
|
||||
expect(disconnect).toHaveBeenCalledTimes(1);
|
||||
expect(dispose).toHaveBeenCalledTimes(1);
|
||||
expect(mockUpdateMCPServerTools).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
tools,
|
||||
|
|
@ -261,9 +327,8 @@ describe('reinitMCPServer — runtime BODY placeholder pre-check (issue #14074)'
|
|||
});
|
||||
|
||||
it('connects normally when the request body provides the placeholder fields', async () => {
|
||||
const disconnect = jest.fn().mockResolvedValue(undefined);
|
||||
mockGetConnection.mockResolvedValue({
|
||||
disconnect,
|
||||
dispose: jest.fn().mockResolvedValue(undefined),
|
||||
fetchTools: jest.fn().mockResolvedValue([]),
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ jest.mock('~/server/services/Config', () => ({
|
|||
setCachedTools: jest.fn(),
|
||||
getCachedTools: jest.fn(),
|
||||
getMCPServerTools: jest.fn(),
|
||||
cacheMCPServerTools: jest.fn(),
|
||||
loadCustomConfig: jest.fn(),
|
||||
}));
|
||||
|
||||
|
|
@ -32,6 +33,7 @@ jest.mock('@librechat/api', () => ({
|
|||
isMCPDomainAllowed: jest.fn(),
|
||||
GenerationJobManager: jest.fn(),
|
||||
buildOAuthToolCallName: jest.fn((name) => name),
|
||||
getUserMCPAuthMap: jest.fn(),
|
||||
/** Mirrors the real resolver so these tests still exercise the wrapper's own
|
||||
* plumbing - loading the request config and degrading on failure - rather than
|
||||
* the resolution logic, which is unit-tested in packages/api. Like the real
|
||||
|
|
@ -53,6 +55,7 @@ jest.mock('~/models', () => ({
|
|||
findToken: jest.fn(),
|
||||
createToken: jest.fn(),
|
||||
updateToken: jest.fn(),
|
||||
findPluginAuthsByKeys: jest.fn(),
|
||||
}));
|
||||
jest.mock('~/server/services/GraphTokenService', () => ({
|
||||
getGraphApiToken: jest.fn(),
|
||||
|
|
@ -69,11 +72,18 @@ jest.mock('~/server/services/Tools/mcp', () => ({
|
|||
|
||||
const { Constants } = require('librechat-data-provider');
|
||||
|
||||
const { getAppConfig } = require('~/server/services/Config');
|
||||
const {
|
||||
getAppConfig,
|
||||
getCachedTools,
|
||||
getMCPServerTools,
|
||||
cacheMCPServerTools,
|
||||
} = require('~/server/services/Config');
|
||||
const { reinitMCPServer } = require('~/server/services/Tools/mcp');
|
||||
const { getUserMCPAuthMap } = require('@librechat/api');
|
||||
const {
|
||||
createMCPTool,
|
||||
healMcpToolNames,
|
||||
getAssistantToolDefinitions,
|
||||
resolveConfigServers,
|
||||
resolveMcpConfigNames,
|
||||
resolveAllMcpConfigs,
|
||||
|
|
@ -81,6 +91,105 @@ const {
|
|||
resolveCollisionAuditNames,
|
||||
} = require('../MCP');
|
||||
|
||||
describe('getAssistantToolDefinitions', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
require('~/config').getMCPManager.mockReset();
|
||||
});
|
||||
|
||||
const req = { user: { id: 'u1', role: 'user' } };
|
||||
const serverConfig = { type: 'streamable-http', url: 'https://app.example.com/mcp' };
|
||||
const toolKey = `search${Constants.mcp_delimiter}app-server`;
|
||||
const mcpDefinition = { type: 'function', function: { name: toolKey } };
|
||||
|
||||
it('combines static definitions with referenced configuration-addressed MCP slices', async () => {
|
||||
getCachedTools.mockResolvedValue({ code_interpreter: { type: 'code_interpreter' } });
|
||||
getAppConfig.mockResolvedValue({ mcpConfig: {} });
|
||||
mockRegistry.ensureConfigServers.mockResolvedValue({});
|
||||
mockRegistry.getAllServerConfigs.mockResolvedValue({ 'app-server': serverConfig });
|
||||
getMCPServerTools.mockResolvedValue({ [toolKey]: mcpDefinition });
|
||||
|
||||
const definitions = await getAssistantToolDefinitions({
|
||||
req,
|
||||
tools: ['code_interpreter', toolKey],
|
||||
});
|
||||
|
||||
expect(definitions).toEqual({
|
||||
code_interpreter: { type: 'code_interpreter' },
|
||||
[toolKey]: mcpDefinition,
|
||||
});
|
||||
expect(getMCPServerTools).toHaveBeenCalledWith('u1', 'app-server', serverConfig);
|
||||
});
|
||||
|
||||
it('recovers and re-caches a referenced server when its slice is missing', async () => {
|
||||
getCachedTools.mockResolvedValue({});
|
||||
getAppConfig.mockResolvedValue({ mcpConfig: {} });
|
||||
mockRegistry.ensureConfigServers.mockResolvedValue({});
|
||||
mockRegistry.getAllServerConfigs.mockResolvedValue({ 'app-server': serverConfig });
|
||||
getMCPServerTools.mockResolvedValue(null);
|
||||
cacheMCPServerTools.mockResolvedValue(undefined);
|
||||
const getServerToolFunctionsSnapshot = jest.fn().mockResolvedValue({
|
||||
tools: { [toolKey]: mcpDefinition },
|
||||
publicationGeneration: 'connection-generation',
|
||||
});
|
||||
require('~/config').getMCPManager.mockReturnValue({ getServerToolFunctionsSnapshot });
|
||||
|
||||
await expect(getAssistantToolDefinitions({ req, tools: [toolKey] })).resolves.toEqual({
|
||||
[toolKey]: mcpDefinition,
|
||||
});
|
||||
expect(cacheMCPServerTools).toHaveBeenCalledWith({
|
||||
userId: 'u1',
|
||||
serverName: 'app-server',
|
||||
serverTools: { [toolKey]: mcpDefinition },
|
||||
serverConfig,
|
||||
publicationGeneration: 'connection-generation',
|
||||
});
|
||||
});
|
||||
|
||||
it('reinitializes a referenced server when its cache and local snapshot are missing', async () => {
|
||||
getCachedTools.mockResolvedValue({});
|
||||
getAppConfig.mockResolvedValue({ mcpConfig: {} });
|
||||
mockRegistry.ensureConfigServers.mockResolvedValue({});
|
||||
mockRegistry.getAllServerConfigs.mockResolvedValue({ 'app-server': serverConfig });
|
||||
getMCPServerTools.mockResolvedValue(null);
|
||||
const getServerToolFunctionsSnapshot = jest.fn().mockResolvedValue({ tools: null });
|
||||
require('~/config').getMCPManager.mockReturnValue({ getServerToolFunctionsSnapshot });
|
||||
const userMCPAuthMap = { 'mcp_app-server': { API_KEY: 'saved' } };
|
||||
getUserMCPAuthMap.mockResolvedValue(userMCPAuthMap);
|
||||
reinitMCPServer.mockResolvedValue({ availableTools: { [toolKey]: mcpDefinition } });
|
||||
|
||||
await expect(getAssistantToolDefinitions({ req, tools: [toolKey] })).resolves.toEqual({
|
||||
[toolKey]: mcpDefinition,
|
||||
});
|
||||
expect(reinitMCPServer).toHaveBeenCalledWith({
|
||||
user: req.user,
|
||||
serverName: 'app-server',
|
||||
serverConfig,
|
||||
userMCPAuthMap,
|
||||
});
|
||||
expect(getUserMCPAuthMap).toHaveBeenCalledWith({
|
||||
userId: 'u1',
|
||||
servers: ['app-server'],
|
||||
findPluginAuthsByKeys: expect.any(Function),
|
||||
});
|
||||
});
|
||||
|
||||
it('propagates config-server resolution failures through the assistant write bridge', async () => {
|
||||
const resolutionError = new Error('config resolution failed');
|
||||
getCachedTools.mockResolvedValue({});
|
||||
getAppConfig.mockResolvedValue({
|
||||
mcpConfig: { 'app-server': { type: 'streamable-http', url: 'https://example.com/mcp' } },
|
||||
});
|
||||
mockRegistry.ensureConfigServers.mockRejectedValue(resolutionError);
|
||||
|
||||
await expect(getAssistantToolDefinitions({ req, tools: [toolKey] })).rejects.toBe(
|
||||
resolutionError,
|
||||
);
|
||||
expect(mockRegistry.getAllServerConfigs).not.toHaveBeenCalled();
|
||||
expect(getMCPServerTools).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('resolveConfigServers', () => {
|
||||
beforeEach(() => jest.clearAllMocks());
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,19 @@
|
|||
const mongoose = require('mongoose');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { mergeAppTools, getAppConfig } = require('./Config');
|
||||
const {
|
||||
registerShutdownTask,
|
||||
setMCPToolsChangedHandler,
|
||||
setMCPToolsChangedGenerationHandler,
|
||||
setMCPToolsChangedGenerationRenewalHandler,
|
||||
setMCPToolsChangedRevisionHandler,
|
||||
} = require('@librechat/api');
|
||||
const { syncStaticTools, mergeAppTools, getAppConfig } = require('./Config');
|
||||
const {
|
||||
getMCPToolsCacheGeneration,
|
||||
renewMCPToolsCacheGeneration,
|
||||
getNextAppToolsPublicationRevision,
|
||||
updateMCPServerTools,
|
||||
} = require('./Config/mcp');
|
||||
const { createMCPServersRegistry, createMCPManager } = require('~/config');
|
||||
|
||||
/**
|
||||
|
|
@ -18,6 +31,36 @@ async function resolveMCPAllowlists(ctx) {
|
|||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Refreshes one server's tools after it reported `notifications/tools/list_changed`.
|
||||
*
|
||||
* A server that builds tools at runtime is the case this exists for: without it the tool list
|
||||
* stayed frozen at connection time and only a restart picked up the change (#7117). The list is
|
||||
* re-fetched from the live connection and written over that server's cache entry, so tools that
|
||||
* disappeared stop being advertised too.
|
||||
*/
|
||||
async function refreshChangedServerTools({
|
||||
serverName,
|
||||
userId,
|
||||
tools,
|
||||
serverConfig,
|
||||
publicationGeneration,
|
||||
publicationRevision,
|
||||
}) {
|
||||
await updateMCPServerTools({
|
||||
userId,
|
||||
serverName,
|
||||
tools,
|
||||
serverConfig,
|
||||
...(publicationGeneration && { publicationGeneration }),
|
||||
...(publicationRevision && { publicationRevision }),
|
||||
});
|
||||
const toolCount = tools.length;
|
||||
logger.info(
|
||||
`[MCP][${serverName}] Tool list changed; refreshed ${toolCount} ${toolCount === 1 ? 'tool' : 'tools'}${userId ? ` for user ${userId}` : ''}`,
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize MCP servers
|
||||
*/
|
||||
|
|
@ -39,16 +82,28 @@ async function initializeMCPs() {
|
|||
|
||||
try {
|
||||
const mcpManager = await createMCPManager(mcpServers || {});
|
||||
setMCPToolsChangedHandler(refreshChangedServerTools);
|
||||
setMCPToolsChangedGenerationHandler(getMCPToolsCacheGeneration);
|
||||
setMCPToolsChangedGenerationRenewalHandler(renewMCPToolsCacheGeneration);
|
||||
setMCPToolsChangedRevisionHandler(({ serverName, configGeneration }) =>
|
||||
getNextAppToolsPublicationRevision(serverName, configGeneration),
|
||||
);
|
||||
registerShutdownTask('MCP app connections', () => mcpManager.disconnectAppServers());
|
||||
|
||||
if (mcpServers && Object.keys(mcpServers).length > 0) {
|
||||
const mcpTools = (await mcpManager.getAppToolFunctions()) || {};
|
||||
await mergeAppTools(mcpTools);
|
||||
try {
|
||||
await mergeAppTools(mcpTools, appConfig.availableTools || {});
|
||||
} finally {
|
||||
await mcpManager.connectAppServers();
|
||||
}
|
||||
const serverCount = Object.keys(mcpServers).length;
|
||||
const toolCount = Object.keys(mcpTools).length;
|
||||
logger.info(
|
||||
`[MCP] Initialized with ${serverCount} configured ${serverCount === 1 ? 'server' : 'servers'} and ${toolCount} ${toolCount === 1 ? 'tool' : 'tools'}.`,
|
||||
);
|
||||
} else {
|
||||
await syncStaticTools(appConfig.availableTools || {});
|
||||
logger.debug('[MCP] No servers configured. MCPManager ready for UI-based servers.');
|
||||
}
|
||||
} catch (error) {
|
||||
|
|
@ -58,3 +113,4 @@ async function initializeMCPs() {
|
|||
}
|
||||
|
||||
module.exports = initializeMCPs;
|
||||
module.exports.refreshChangedServerTools = refreshChangedServerTools;
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ jest.mock('@librechat/data-schemas', () => ({
|
|||
|
||||
// Mock config functions
|
||||
const mockGetAppConfig = jest.fn();
|
||||
const mockSyncStaticTools = jest.fn();
|
||||
const mockMergeAppTools = jest.fn();
|
||||
|
||||
jest.mock('./Config', () => ({
|
||||
|
|
@ -36,12 +37,17 @@ jest.mock('./Config', () => ({
|
|||
get mergeAppTools() {
|
||||
return mockMergeAppTools;
|
||||
},
|
||||
get syncStaticTools() {
|
||||
return mockSyncStaticTools;
|
||||
},
|
||||
}));
|
||||
|
||||
// Mock MCP singletons
|
||||
const mockCreateMCPServersRegistry = jest.fn();
|
||||
const mockCreateMCPManager = jest.fn();
|
||||
const mockMCPManagerInstance = {
|
||||
connectAppServers: jest.fn(),
|
||||
disconnectAppServers: jest.fn(),
|
||||
getAppToolFunctions: jest.fn(),
|
||||
};
|
||||
|
||||
|
|
@ -54,6 +60,49 @@ jest.mock('~/config', () => ({
|
|||
},
|
||||
}));
|
||||
|
||||
const mockSetMCPToolsChangedHandler = jest.fn();
|
||||
const mockSetMCPToolsChangedGenerationHandler = jest.fn();
|
||||
const mockSetMCPToolsChangedGenerationRenewalHandler = jest.fn();
|
||||
const mockSetMCPToolsChangedRevisionHandler = jest.fn();
|
||||
const mockRegisterShutdownTask = jest.fn();
|
||||
const mockUpdateMCPServerTools = jest.fn();
|
||||
const mockGetMCPToolsCacheGeneration = jest.fn();
|
||||
const mockRenewMCPToolsCacheGeneration = jest.fn();
|
||||
const mockGetNextAppToolsPublicationRevision = jest.fn();
|
||||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
get registerShutdownTask() {
|
||||
return mockRegisterShutdownTask;
|
||||
},
|
||||
get setMCPToolsChangedHandler() {
|
||||
return mockSetMCPToolsChangedHandler;
|
||||
},
|
||||
get setMCPToolsChangedGenerationHandler() {
|
||||
return mockSetMCPToolsChangedGenerationHandler;
|
||||
},
|
||||
get setMCPToolsChangedGenerationRenewalHandler() {
|
||||
return mockSetMCPToolsChangedGenerationRenewalHandler;
|
||||
},
|
||||
get setMCPToolsChangedRevisionHandler() {
|
||||
return mockSetMCPToolsChangedRevisionHandler;
|
||||
},
|
||||
}));
|
||||
|
||||
jest.mock('./Config/mcp', () => ({
|
||||
get updateMCPServerTools() {
|
||||
return mockUpdateMCPServerTools;
|
||||
},
|
||||
get getMCPToolsCacheGeneration() {
|
||||
return mockGetMCPToolsCacheGeneration;
|
||||
},
|
||||
get renewMCPToolsCacheGeneration() {
|
||||
return mockRenewMCPToolsCacheGeneration;
|
||||
},
|
||||
get getNextAppToolsPublicationRevision() {
|
||||
return mockGetNextAppToolsPublicationRevision;
|
||||
},
|
||||
}));
|
||||
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const initializeMCPs = require('./initializeMCPs');
|
||||
|
||||
|
|
@ -65,6 +114,9 @@ describe('initializeMCPs', () => {
|
|||
mockCreateMCPServersRegistry.mockReturnValue(undefined);
|
||||
mockCreateMCPManager.mockResolvedValue(mockMCPManagerInstance);
|
||||
mockMCPManagerInstance.getAppToolFunctions.mockResolvedValue({});
|
||||
mockMCPManagerInstance.connectAppServers.mockResolvedValue(undefined);
|
||||
mockMCPManagerInstance.disconnectAppServers.mockResolvedValue(undefined);
|
||||
mockSyncStaticTools.mockResolvedValue(undefined);
|
||||
mockMergeAppTools.mockResolvedValue(undefined);
|
||||
});
|
||||
|
||||
|
|
@ -183,6 +235,36 @@ describe('initializeMCPs', () => {
|
|||
expect(mockCreateMCPManager).toHaveBeenCalledWith(mcpServers);
|
||||
});
|
||||
|
||||
it('should register app connections for graceful shutdown', async () => {
|
||||
mockGetAppConfig.mockResolvedValue({ mcpConfig: null });
|
||||
|
||||
await initializeMCPs();
|
||||
|
||||
expect(mockRegisterShutdownTask).toHaveBeenCalledWith(
|
||||
'MCP app connections',
|
||||
expect.any(Function),
|
||||
);
|
||||
const shutdown = mockRegisterShutdownTask.mock.calls[0][1];
|
||||
await shutdown();
|
||||
expect(mockMCPManagerInstance.disconnectAppServers).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should wire app publication revision allocation into the cache store', async () => {
|
||||
mockGetAppConfig.mockResolvedValue({ mcpConfig: null });
|
||||
mockGetNextAppToolsPublicationRevision.mockResolvedValue('9');
|
||||
|
||||
await initializeMCPs();
|
||||
|
||||
const allocateRevision = mockSetMCPToolsChangedRevisionHandler.mock.calls[0][0];
|
||||
await expect(
|
||||
allocateRevision({ serverName: 'dynamic', configGeneration: 'config-generation' }),
|
||||
).resolves.toBe('9');
|
||||
expect(mockGetNextAppToolsPublicationRevision).toHaveBeenCalledWith(
|
||||
'dynamic',
|
||||
'config-generation',
|
||||
);
|
||||
});
|
||||
|
||||
it('should throw and log error if MCPManager initialization fails', async () => {
|
||||
const managerError = new Error('Manager initialization failed');
|
||||
mockCreateMCPManager.mockRejectedValue(managerError);
|
||||
|
|
@ -197,21 +279,24 @@ describe('initializeMCPs', () => {
|
|||
});
|
||||
|
||||
describe('Tool merging behavior', () => {
|
||||
it('should NOT merge tools when no configured servers exist', async () => {
|
||||
it('should skip app catalog discovery when no configured servers exist', async () => {
|
||||
mockGetAppConfig.mockResolvedValue({
|
||||
mcpConfig: null, // No configured servers
|
||||
availableTools: { builtin: { type: 'function' } },
|
||||
});
|
||||
|
||||
await initializeMCPs();
|
||||
|
||||
expect(mockMCPManagerInstance.getAppToolFunctions).not.toHaveBeenCalled();
|
||||
expect(mockMergeAppTools).not.toHaveBeenCalled();
|
||||
expect(mockSyncStaticTools).toHaveBeenCalledWith({ builtin: { type: 'function' } });
|
||||
expect(mockMCPManagerInstance.connectAppServers).not.toHaveBeenCalled();
|
||||
expect(logger.debug).toHaveBeenCalledWith(
|
||||
'[MCP] No servers configured. MCPManager ready for UI-based servers.',
|
||||
);
|
||||
});
|
||||
|
||||
it('should NOT merge tools when mcpConfig is empty object', async () => {
|
||||
it('should skip app catalog discovery when mcpConfig is empty', async () => {
|
||||
mockGetAppConfig.mockResolvedValue({
|
||||
mcpConfig: {}, // Empty object
|
||||
});
|
||||
|
|
@ -220,6 +305,8 @@ describe('initializeMCPs', () => {
|
|||
|
||||
expect(mockMCPManagerInstance.getAppToolFunctions).not.toHaveBeenCalled();
|
||||
expect(mockMergeAppTools).not.toHaveBeenCalled();
|
||||
expect(mockSyncStaticTools).toHaveBeenCalledWith({});
|
||||
expect(mockMCPManagerInstance.connectAppServers).not.toHaveBeenCalled();
|
||||
expect(logger.debug).toHaveBeenCalledWith(
|
||||
'[MCP] No servers configured. MCPManager ready for UI-based servers.',
|
||||
);
|
||||
|
|
@ -239,7 +326,11 @@ describe('initializeMCPs', () => {
|
|||
await initializeMCPs();
|
||||
|
||||
expect(mockMCPManagerInstance.getAppToolFunctions).toHaveBeenCalledTimes(1);
|
||||
expect(mockMergeAppTools).toHaveBeenCalledWith(mcpTools);
|
||||
expect(mockMergeAppTools).toHaveBeenCalledWith(mcpTools, {});
|
||||
expect(mockMCPManagerInstance.connectAppServers).toHaveBeenCalledTimes(1);
|
||||
expect(mockMergeAppTools.mock.invocationCallOrder[0]).toBeLessThan(
|
||||
mockMCPManagerInstance.connectAppServers.mock.invocationCallOrder[0],
|
||||
);
|
||||
expect(logger.info).toHaveBeenCalledWith(
|
||||
'[MCP] Initialized with 1 configured server and 2 tools.',
|
||||
);
|
||||
|
|
@ -253,11 +344,21 @@ describe('initializeMCPs', () => {
|
|||
await initializeMCPs();
|
||||
|
||||
// Should use empty object fallback
|
||||
expect(mockMergeAppTools).toHaveBeenCalledWith({});
|
||||
expect(mockMergeAppTools).toHaveBeenCalledWith({}, {});
|
||||
expect(logger.info).toHaveBeenCalledWith(
|
||||
'[MCP] Initialized with 1 configured server and 0 tools.',
|
||||
);
|
||||
});
|
||||
|
||||
it('should connect app servers when startup cache synchronization fails', async () => {
|
||||
const mcpServers = { 'test-server': { type: 'sse', url: 'http://localhost:3001' } };
|
||||
mockGetAppConfig.mockResolvedValue({ mcpConfig: mcpServers });
|
||||
mockMergeAppTools.mockRejectedValueOnce(new Error('cache lock timed out'));
|
||||
|
||||
await expect(initializeMCPs()).rejects.toThrow('cache lock timed out');
|
||||
|
||||
expect(mockMCPManagerInstance.connectAppServers).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Initialization order', () => {
|
||||
|
|
@ -315,3 +416,45 @@ describe('initializeMCPs', () => {
|
|||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('refreshChangedServerTools', () => {
|
||||
const { refreshChangedServerTools } = require('./initializeMCPs');
|
||||
const event = {
|
||||
serverName: 'dynamic',
|
||||
serverConfig: { type: 'streamable-http', url: 'https://mcp.example.com' },
|
||||
tools: [{ name: 'tool', inputSchema: { type: 'object' } }],
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
it('publishes the complete refreshed snapshot in its original cache scope', async () => {
|
||||
await refreshChangedServerTools({ ...event, userId: 'user-1' });
|
||||
|
||||
expect(mockUpdateMCPServerTools).toHaveBeenCalledWith({ ...event, userId: 'user-1' });
|
||||
expect(logger.info).toHaveBeenCalledWith(
|
||||
'[MCP][dynamic] Tool list changed; refreshed 1 tool for user user-1',
|
||||
);
|
||||
});
|
||||
|
||||
it('publishes an empty app-level snapshot so removals take effect', async () => {
|
||||
await refreshChangedServerTools({ ...event, tools: [] });
|
||||
|
||||
expect(mockUpdateMCPServerTools).toHaveBeenCalledWith({ ...event, tools: [] });
|
||||
});
|
||||
|
||||
it('is registered as the tools-changed handler during initialization', async () => {
|
||||
mockGetAppConfig.mockResolvedValue({ mcpConfig: null, mcpSettings: {} });
|
||||
|
||||
await initializeMCPs();
|
||||
|
||||
expect(mockSetMCPToolsChangedHandler).toHaveBeenCalledWith(refreshChangedServerTools);
|
||||
expect(mockSetMCPToolsChangedGenerationHandler).toHaveBeenCalledWith(
|
||||
mockGetMCPToolsCacheGeneration,
|
||||
);
|
||||
expect(mockSetMCPToolsChangedGenerationRenewalHandler).toHaveBeenCalledWith(
|
||||
mockRenewMCPToolsCacheGeneration,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue