Add include_tools and exclude_tools filtering for MCP servers

ee/vellum_ee/workflows/display/nodes/vellum/tests/test_tool_calling_node.py

@@ -246,6 +246,16 @@ class Workflow(BaseWorkflow):
                                 "environment_variable": "my-api-key-header-value",
                             },
                         },
+                        {
+                            "id": "0ce4327e-72ab-4fb3-83ab-9cce5d398388",
+                            "key": "include_tools",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
+                        {
+                            "id": "9ee4fd57-b73a-4b97-991e-14cad22645fa",
+                            "key": "exclude_tools",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
                     ],
                     "definition": {"name": "MCPServer", "module": ["vellum", "workflows", "types", "definition"]},
                 }
@@ -254,7 +264,60 @@ class Workflow(BaseWorkflow):
     }
 
 
+def test_serialize_node__tool_calling_node__mcp_server_with_tool_filtering():
+    """Tests that MCPServer with include_tools and exclude_tools serializes correctly."""
+
+    # GIVEN a tool calling node with an mcp server that has tool filtering configured
+    class MyToolCallingNode(ToolCallingNode):
+        functions = [
+            MCPServer(
+                name="my-mcp-server",
+                url="https://my-mcp-server.com",
+                authorization_type=AuthorizationType.API_KEY,
+                api_key_header_key="my-api-key-header-key",
+                api_key_header_value=EnvironmentVariableReference(name="my-api-key-header-value"),
+                include_tools=["tool-a", "tool-b"],
+                exclude_tools=["tool-c"],
+            )
+        ]
+
+    # AND a workflow with the tool calling node
+    class Workflow(BaseWorkflow):
+        graph = MyToolCallingNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the mcp server with tool filtering
+    my_tool_calling_node = next(
+        node
+        for node in serialized_workflow["workflow_raw_data"]["nodes"]
+        if node["id"] == str(MyToolCallingNode.__id__)
+    )
+
+    functions_attribute = next(
+        attribute for attribute in my_tool_calling_node["attributes"] if attribute["name"] == "functions"
+    )
+
+    # AND the include_tools and exclude_tools should be serialized as JSON arrays
+    entries = functions_attribute["value"]["items"][0]["entries"]
+    include_tools_entry = next(entry for entry in entries if entry["key"] == "include_tools")
+    exclude_tools_entry = next(entry for entry in entries if entry["key"] == "exclude_tools")
+
+    assert include_tools_entry["value"] == {
+        "type": "CONSTANT_VALUE",
+        "value": {"type": "JSON", "value": ["tool-a", "tool-b"]},
+    }
+    assert exclude_tools_entry["value"] == {
+        "type": "CONSTANT_VALUE",
+        "value": {"type": "JSON", "value": ["tool-c"]},
+    }
+
+
 def test_serialize_node__tool_calling_node__mcp_server_no_authorization():
+    """Tests that MCPServer without authorization serializes correctly."""
+
     # GIVEN a tool calling node with an mcp server
     class MyToolCallingNode(ToolCallingNode):
         functions = [
@@ -300,6 +363,8 @@ class Workflow(BaseWorkflow):
                         "bearer_token_value": None,
                         "api_key_header_key": None,
                         "api_key_header_value": None,
+                        "include_tools": None,
+                        "exclude_tools": None,
                     }
                 ],
             },

ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_mcp_serialization.py

@@ -96,6 +96,16 @@ def test_serialize_workflow():
                             "key": "api_key_header_value",
                             "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
                         },
+                        {
+                            "id": "90981dfe-9dc8-4f40-a5d6-b8b42edfc075",
+                            "key": "include_tools",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
+                        {
+                            "id": "d48a1658-0d94-49fc-9499-5f400079eebd",
+                            "key": "exclude_tools",
+                            "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+                        },
                     ],
                     "definition": {"name": "MCPServer", "module": ["vellum", "workflows", "types", "definition"]},
                 }

src/vellum/workflows/integrations/mcp_service.py

@@ -264,18 +264,36 @@ def execute_tool(self, tool_def: MCPToolDefinition, arguments: Dict[str, Any]) -
             )
 
     def hydrate_tool_definitions(self, server_def: MCPServer) -> List[MCPToolDefinition]:
-        """Hydrate an MCPToolDefinition with detailed information from the MCP server."""
+        """Hydrate an MCPToolDefinition with detailed information from the MCP server.
+
+        Tool filtering is applied based on include_tools and exclude_tools:
+        - include_tools acts as a whitelist: only tools in this list are included
+        - exclude_tools acts as a blacklist: tools in this list are filtered out
+        - exclude_tools takes precedence: if a tool appears in both lists, it is excluded
+        """
         try:
             tools = self.list_tools(server_def)
 
+            filtered_tools = []
+            for tool in tools:
+                tool_name = tool["name"]
+
+                if server_def.include_tools is not None and tool_name not in server_def.include_tools:
+                    continue
+
+                if server_def.exclude_tools is not None and tool_name in server_def.exclude_tools:
+                    continue
+
+                filtered_tools.append(tool)
+
             return [
                 MCPToolDefinition(
                     name=tool["name"],
                     server=server_def,
                     description=tool["description"],
                     parameters=tool["inputSchema"],
                 )
-                for tool in tools
+                for tool in filtered_tools
             ]
         except Exception as e:
             logger.warning(f"Failed to hydrate MCP server '{server_def.name}': {e}")

src/vellum/workflows/integrations/tests/test_mcp_service.py

@@ -271,3 +271,125 @@ def test_mcp_service_call_tool_includes_stacktrace_and_raw_data_on_error():
         assert exc_info.value.raw_data["operation"] == "call_tool"
         assert exc_info.value.raw_data["error_type"] == "RuntimeError"
         assert exc_info.value.raw_data["error_message"] == "Tool execution failed"
+
+
+def test_mcp_service_hydrate_tool_definitions__include_tools_filters_to_whitelist():
+    """
+    Tests that include_tools acts as a whitelist, only including specified tools.
+    """
+    # GIVEN an MCP server with include_tools specified
+    sample_mcp_server = MCPServer(
+        name="test-server",
+        url="https://test.mcp.server.com/mcp",
+        include_tools=["tool-a", "tool-c"],
+    )
+
+    # AND a mock MCP service that returns multiple tools
+    with mock.patch("vellum.workflows.integrations.mcp_service.asyncio.run") as mock_run:
+        mock_run.return_value = [
+            {"name": "tool-a", "description": "Tool A", "inputSchema": {}},
+            {"name": "tool-b", "description": "Tool B", "inputSchema": {}},
+            {"name": "tool-c", "description": "Tool C", "inputSchema": {}},
+        ]
+
+        # WHEN we hydrate tool definitions
+        service = MCPService()
+        tool_definitions = service.hydrate_tool_definitions(sample_mcp_server)
+
+        # THEN only the tools in include_tools should be returned
+        assert len(tool_definitions) == 2
+        tool_names = [t.name for t in tool_definitions]
+        assert "tool-a" in tool_names
+        assert "tool-c" in tool_names
+        assert "tool-b" not in tool_names
+
+
+def test_mcp_service_hydrate_tool_definitions__exclude_tools_filters_out_blacklist():
+    """
+    Tests that exclude_tools acts as a blacklist, filtering out specified tools.
+    """
+    # GIVEN an MCP server with exclude_tools specified
+    sample_mcp_server = MCPServer(
+        name="test-server",
+        url="https://test.mcp.server.com/mcp",
+        exclude_tools=["tool-b"],
+    )
+
+    # AND a mock MCP service that returns multiple tools
+    with mock.patch("vellum.workflows.integrations.mcp_service.asyncio.run") as mock_run:
+        mock_run.return_value = [
+            {"name": "tool-a", "description": "Tool A", "inputSchema": {}},
+            {"name": "tool-b", "description": "Tool B", "inputSchema": {}},
+            {"name": "tool-c", "description": "Tool C", "inputSchema": {}},
+        ]
+
+        # WHEN we hydrate tool definitions
+        service = MCPService()
+        tool_definitions = service.hydrate_tool_definitions(sample_mcp_server)
+
+        # THEN all tools except the excluded one should be returned
+        assert len(tool_definitions) == 2
+        tool_names = [t.name for t in tool_definitions]
+        assert "tool-a" in tool_names
+        assert "tool-c" in tool_names
+        assert "tool-b" not in tool_names
+
+
+def test_mcp_service_hydrate_tool_definitions__exclude_takes_precedence_over_include():
+    """
+    Tests that exclude_tools takes precedence when a tool appears in both lists.
+    """
+    # GIVEN an MCP server with both include_tools and exclude_tools
+    # AND tool-a appears in both lists
+    sample_mcp_server = MCPServer(
+        name="test-server",
+        url="https://test.mcp.server.com/mcp",
+        include_tools=["tool-a", "tool-b"],
+        exclude_tools=["tool-a"],
+    )
+
+    # AND a mock MCP service that returns multiple tools
+    with mock.patch("vellum.workflows.integrations.mcp_service.asyncio.run") as mock_run:
+        mock_run.return_value = [
+            {"name": "tool-a", "description": "Tool A", "inputSchema": {}},
+            {"name": "tool-b", "description": "Tool B", "inputSchema": {}},
+            {"name": "tool-c", "description": "Tool C", "inputSchema": {}},
+        ]
+
+        # WHEN we hydrate tool definitions
+        service = MCPService()
+        tool_definitions = service.hydrate_tool_definitions(sample_mcp_server)
+
+        # THEN tool-a should be excluded even though it's in include_tools
+        assert len(tool_definitions) == 1
+        assert tool_definitions[0].name == "tool-b"
+
+
+def test_mcp_service_hydrate_tool_definitions__no_filtering_when_both_none():
+    """
+    Tests that all tools are returned when neither include_tools nor exclude_tools is specified.
+    """
+    # GIVEN an MCP server with no filtering configured
+    sample_mcp_server = MCPServer(
+        name="test-server",
+        url="https://test.mcp.server.com/mcp",
+    )
+
+    # AND a mock MCP service that returns multiple tools
+    with mock.patch("vellum.workflows.integrations.mcp_service.asyncio.run") as mock_run:
+        mock_run.return_value = [
+            {"name": "tool-a", "description": "Tool A", "inputSchema": {}},
+            {"name": "tool-b", "description": "Tool B", "inputSchema": {}},
+            {"name": "tool-c", "description": "Tool C", "inputSchema": {}},
+        ]
+
+        # WHEN we hydrate tool definitions
+        service = MCPService()
+        tool_definitions = service.hydrate_tool_definitions(sample_mcp_server)
+
+        # THEN all tools should be returned
+        assert len(tool_definitions) == 3
+        tool_names = [t.name for t in tool_definitions]
+        assert "tool-a" in tool_names
+        assert "tool-b" in tool_names
+        assert "tool-c" in tool_names

src/vellum/workflows/types/definition.py

@@ -205,6 +205,8 @@ class MCPServer(UniversalBaseModel):
     bearer_token_value: Optional[Union[str, EnvironmentVariableReference]] = None
     api_key_header_key: Optional[str] = None
     api_key_header_value: Optional[Union[str, EnvironmentVariableReference]] = None
+    include_tools: Optional[List[str]] = None
+    exclude_tools: Optional[List[str]] = None
 
     model_config = {"arbitrary_types_allowed": True}