diff options
Diffstat (limited to 'packages/core/src/tools/mcp-client.test.ts')
| -rw-r--r-- | packages/core/src/tools/mcp-client.test.ts | 109 |
1 files changed, 109 insertions, 0 deletions
diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts index 09f7951b..1413a4f8 100644 --- a/packages/core/src/tools/mcp-client.test.ts +++ b/packages/core/src/tools/mcp-client.test.ts @@ -668,6 +668,115 @@ describe('discoverMcpTools', () => { clientInstances[clientInstances.length - 1]?.value; 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', () => { |
