diff --git a/app/controllers/concerns/foreman_wds/discovered_hosts_controller_extensions.rb b/app/controllers/concerns/foreman_wds/discovered_hosts_controller_extensions.rb
new file mode 100644
index 0000000000000000000000000000000000000000..61ddb46e3c9429858dd381e3c4b29fe2a4fcdbcf
--- /dev/null
+++ b/app/controllers/concerns/foreman_wds/discovered_hosts_controller_extensions.rb
@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module ForemanWds
+  module DiscoveredHostsControllerExtensions
+    def action_permission
+      return :edit if params[:action] == 'wds_server_selected'
+
+      super
+    end
+  end
+end
diff --git a/app/controllers/concerns/foreman_wds/hosts_controller_extensions.rb b/app/controllers/concerns/foreman_wds/hosts_controller_extensions.rb
index 9ddf276713b25c23d9300ca2073e70b6c1af141b..b2af16f79b16ca820d592f4ed44e014135210525 100644
--- a/app/controllers/concerns/foreman_wds/hosts_controller_extensions.rb
+++ b/app/controllers/concerns/foreman_wds/hosts_controller_extensions.rb
@@ -7,7 +7,7 @@ module ForemanWds
     included do
       before_action :cleanup_wds_params
 
-      define_action_permission %w[wds_server_selected], :edit
+      define_action_permission %w[wds_server_selected], :edit if respond_to? :define_action_permission
     end
 
     def wds_server_selected
diff --git a/lib/foreman_wds/engine.rb b/lib/foreman_wds/engine.rb
index e7499ec72b1c9578b2ab24c4325fdc183cc3fec0..336c942f21ee44dedb388d72dacb8a7afda34390 100644
--- a/lib/foreman_wds/engine.rb
+++ b/lib/foreman_wds/engine.rb
@@ -33,6 +33,7 @@ module ForemanWds
 
       if Foreman::Plugin.installed?('foreman_discovery')
         DiscoveredHostsController.include ForemanWds::HostsControllerExtensions
+        DiscoveredHostsController.prepend ForemanWds::DiscoveredHostsControllerExtensions
       end
     rescue StandardError => e
       Rails.logger.fatal "foreman_wds: skipping engine hook (#{e})"