Add excludeTools and includeTools to mcpServers config (#2976)
This commit is contained in:
parent
aa8e5776eb
commit
355fb4ac67
|
@ -84,6 +84,8 @@ export class MCPServerConfig {
|
||||||
readonly trust?: boolean,
|
readonly trust?: boolean,
|
||||||
// Metadata
|
// Metadata
|
||||||
readonly description?: string,
|
readonly description?: string,
|
||||||
|
readonly includeTools?: string[],
|
||||||
|
readonly excludeTools?: string[],
|
||||||
) {}
|
) {}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
@ -668,6 +668,115 @@ describe('discoverMcpTools', () => {
|
||||||
clientInstances[clientInstances.length - 1]?.value;
|
clientInstances[clientInstances.length - 1]?.value;
|
||||||
expect(lastClientInstance?.onerror).toEqual(expect.any(Function));
|
expect(lastClientInstance?.onerror).toEqual(expect.any(Function));
|
||||||
});
|
});
|
||||||
|
|
||||||
|
describe('Tool Filtering', () => {
|
||||||
|
const mockTools = [
|
||||||
|
{
|
||||||
|
name: 'toolA',
|
||||||
|
description: 'descA',
|
||||||
|
inputSchema: { type: 'object' as const, properties: {} },
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: 'toolB',
|
||||||
|
description: 'descB',
|
||||||
|
inputSchema: { type: 'object' as const, properties: {} },
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: 'toolC',
|
||||||
|
description: 'descC',
|
||||||
|
inputSchema: { type: 'object' as const, properties: {} },
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
vi.mocked(Client.prototype.listTools).mockResolvedValue({
|
||||||
|
tools: mockTools,
|
||||||
|
});
|
||||||
|
mockToolRegistry.getToolsByServer.mockReturnValue([
|
||||||
|
expect.any(DiscoveredMCPTool),
|
||||||
|
]);
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should only include specified tools with includeTools', async () => {
|
||||||
|
const serverConfig: MCPServerConfig = {
|
||||||
|
command: './mcp-include',
|
||||||
|
includeTools: ['toolA', 'toolC'],
|
||||||
|
};
|
||||||
|
mockConfig.getMcpServers.mockReturnValue({
|
||||||
|
'include-server': serverConfig,
|
||||||
|
});
|
||||||
|
|
||||||
|
await discoverMcpTools(
|
||||||
|
mockConfig.getMcpServers() ?? {},
|
||||||
|
mockConfig.getMcpServerCommand(),
|
||||||
|
mockToolRegistry as any,
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(mockToolRegistry.registerTool).toHaveBeenCalledTimes(2);
|
||||||
|
expect(mockToolRegistry.registerTool).toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({ serverToolName: 'toolA' }),
|
||||||
|
);
|
||||||
|
expect(mockToolRegistry.registerTool).toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({ serverToolName: 'toolC' }),
|
||||||
|
);
|
||||||
|
expect(mockToolRegistry.registerTool).not.toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({ serverToolName: 'toolB' }),
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should exclude specified tools with excludeTools', async () => {
|
||||||
|
const serverConfig: MCPServerConfig = {
|
||||||
|
command: './mcp-exclude',
|
||||||
|
excludeTools: ['toolB'],
|
||||||
|
};
|
||||||
|
mockConfig.getMcpServers.mockReturnValue({
|
||||||
|
'exclude-server': serverConfig,
|
||||||
|
});
|
||||||
|
|
||||||
|
await discoverMcpTools(
|
||||||
|
mockConfig.getMcpServers() ?? {},
|
||||||
|
mockConfig.getMcpServerCommand(),
|
||||||
|
mockToolRegistry as any,
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(mockToolRegistry.registerTool).toHaveBeenCalledTimes(2);
|
||||||
|
expect(mockToolRegistry.registerTool).toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({ serverToolName: 'toolA' }),
|
||||||
|
);
|
||||||
|
expect(mockToolRegistry.registerTool).toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({ serverToolName: 'toolC' }),
|
||||||
|
);
|
||||||
|
expect(mockToolRegistry.registerTool).not.toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({ serverToolName: 'toolB' }),
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should handle both includeTools and excludeTools', async () => {
|
||||||
|
const serverConfig: MCPServerConfig = {
|
||||||
|
command: './mcp-both',
|
||||||
|
includeTools: ['toolA', 'toolB'],
|
||||||
|
excludeTools: ['toolB'],
|
||||||
|
};
|
||||||
|
mockConfig.getMcpServers.mockReturnValue({ 'both-server': serverConfig });
|
||||||
|
|
||||||
|
await discoverMcpTools(
|
||||||
|
mockConfig.getMcpServers() ?? {},
|
||||||
|
mockConfig.getMcpServerCommand(),
|
||||||
|
mockToolRegistry as any,
|
||||||
|
);
|
||||||
|
|
||||||
|
expect(mockToolRegistry.registerTool).toHaveBeenCalledTimes(1);
|
||||||
|
expect(mockToolRegistry.registerTool).toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({ serverToolName: 'toolA' }),
|
||||||
|
);
|
||||||
|
expect(mockToolRegistry.registerTool).not.toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({ serverToolName: 'toolB' }),
|
||||||
|
);
|
||||||
|
expect(mockToolRegistry.registerTool).not.toHaveBeenCalledWith(
|
||||||
|
expect.objectContaining({ serverToolName: 'toolC' }),
|
||||||
|
);
|
||||||
|
});
|
||||||
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
describe('sanitizeParameters', () => {
|
describe('sanitizeParameters', () => {
|
||||||
|
|
|
@ -305,6 +305,26 @@ async function connectAndDiscover(
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const { includeTools, excludeTools } = mcpServerConfig;
|
||||||
|
const toolName = funcDecl.name;
|
||||||
|
|
||||||
|
let isEnabled = false;
|
||||||
|
if (includeTools === undefined) {
|
||||||
|
isEnabled = true;
|
||||||
|
} else {
|
||||||
|
isEnabled = includeTools.some(
|
||||||
|
(tool) => tool === toolName || tool.startsWith(`${toolName}(`),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (excludeTools?.includes(toolName)) {
|
||||||
|
isEnabled = false;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (!isEnabled) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
let toolNameForModel = funcDecl.name;
|
let toolNameForModel = funcDecl.name;
|
||||||
|
|
||||||
// Replace invalid characters (based on 400 error message from Gemini API) with underscores
|
// Replace invalid characters (based on 400 error message from Gemini API) with underscores
|
||||||
|
|
Loading…
Reference in New Issue