(GH-177) Dynamically unregister capabilities

lib/puppet-languageserver/language_client.rb

@@ -2,15 +2,28 @@
 
 module PuppetLanguageServer
   class LanguageClient
-    def initialize
+    attr_reader :message_router
+
+    def initialize(message_router)
+      @message_router = message_router
       @client_capabilites = {}
+
+      # Internal registry of dynamic registrations and their current state
+      # @registrations[ <[String] method_name>] = [
+      #  {
+      #    :id         => [String] Request ID. Used for de-registration
+      #    :registered => [Boolean] true | false
+      #    :state      => [Enum] :pending | :complete
+      #   }
+      # ]
+      @registrations = {}
     end
 
     def client_capability(*names)
       safe_hash_traverse(@client_capabilites, *names)
     end
 
-    def send_configuration_request(message_router)
+    def send_configuration_request
       params = LSP::ConfigurationParams.new.from_h!('items' => [])
       params.items << LSP::ConfigurationItem.new.from_h!('section' => 'puppet')
 
@@ -30,32 +43,115 @@ def parse_lsp_configuration_settings!(settings = [{}])
       # end
     end
 
-    def register_capability(message_router, method, options = {})
-      id = SecureRandom.uuid
+    def capability_registrations(method)
+      return [{ :registered => false, :state => :complete }] if @registrations[method].nil? || @registrations[method].empty?
+      @registrations[method].dup
+    end
+
+    def register_capability(method, options = {})
+      id = new_request_id
 
       PuppetLanguageServer.log_message(:info, "Attempting to dynamically register the #{method} method with id #{id}")
 
+      if @registrations[method] && @registrations[method].select { |i| i[:state] == :pending }.count > 0
+        # The protocol doesn't specify whether this is allowed and is probably per client specific. For the moment we will allow
+        # the registration to be sent but log a message that something may be wrong.
+        PuppetLanguageServer.log_message(:warn, "A dynamic registration/deregistration for the #{method} method is already in progress")
+      end
+
       params = LSP::RegistrationParams.new.from_h!('registrations' => [])
       params.registrations << LSP::Registration.new.from_h!('id' => id, 'method' => method, 'registerOptions' => options)
+      # Note - Don't put more than one method per request even though you can.  It makes decoding errors much harder!
+
+      @registrations[method] = [] if @registrations[method].nil?
+      @registrations[method] << { :registered => false, :state => :pending, :id => id }
 
       message_router.json_rpc_handler.send_client_request('client/registerCapability', params)
       true
     end
 
-    def parse_register_capability_response!(message_router, _response, original_request)
+    def unregister_capability(method)
+      if @registrations[method].nil?
+        PuppetLanguageServer.log_message(:debug, "No registrations to deregister for the #{method}")
+        return true
+      end
+
+      params = LSP::UnregistrationParams.new.from_h!('unregisterations' => [])
+      @registrations[method].each do |reg|
+        next if reg[:id].nil?
+        PuppetLanguageServer.log_message(:warn, "A dynamic registration/deregistration for the #{method} method, with id #{reg[:id]} is already in progress") if reg[:state] == :pending
+        # Ignore registrations that don't need to be unregistered
+        next if reg[:state] == :complete && !reg[:registered]
+        params.unregisterations << LSP::Unregistration.new.from_h!('id' => reg[:id], 'method' => method)
+        reg[:state] = :pending
+      end
+
+      if params.unregisterations.count.zero?
+        PuppetLanguageServer.log_message(:debug, "Nothing to deregister for the #{method} method")
+        return true
+      end
+
+      message_router.json_rpc_handler.send_client_request('client/unregisterCapability', params)
+      true
+    end
+
+    def parse_register_capability_response!(response, original_request)
       raise 'Response is not from client/registerCapability request' unless original_request['method'] == 'client/registerCapability'
+
+      unless response.key?('result')
+        original_request['params'].registrations.each do |reg|
+          # Mark the registration as completed and failed
+          @registrations[reg.method__lsp] = [] if @registrations[reg.method__lsp].nil?
+          @registrations[reg.method__lsp].select { |i| i[:id] == reg.id }.each { |i| i[:registered] = false; i[:state] = :complete } # rubocop:disable Style/Semicolon This is fine
+        end
+        return true
+      end
+
       original_request['params'].registrations.each do |reg|
         PuppetLanguageServer.log_message(:info, "Succesfully dynamically registered the #{reg.method__lsp} method")
+
+        # Mark the registration as completed and succesful
+        @registrations[reg.method__lsp] = [] if @registrations[reg.method__lsp].nil?
+        @registrations[reg.method__lsp].select { |i| i[:id] == reg.id }.each { |i| i[:registered] = true; i[:state] = :complete } # rubocop:disable Style/Semicolon This is fine
+
         # If we just registered the workspace/didChangeConfiguration method then
         # also trigger a configuration request to get the initial state
-        send_configuration_request(message_router) if reg.method__lsp == 'workspace/didChangeConfiguration'
+        send_configuration_request if reg.method__lsp == 'workspace/didChangeConfiguration'
+      end
+
+      true
+    end
+
+    def parse_unregister_capability_response!(response, original_request)
+      raise 'Response is not from client/unregisterCapability request' unless original_request['method'] == 'client/unregisterCapability'
+
+      unless response.key?('result')
+        original_request['params'].unregisterations.each do |reg|
+          # Mark the registration as completed and failed
+          @registrations[reg.method__lsp] = [] if @registrations[reg.method__lsp].nil?
+          @registrations[reg.method__lsp].select { |i| i[:id] == reg.id && i[:registered] }.each { |i| i[:state] = :complete }
+          @registrations[reg.method__lsp].delete_if { |i| i[:id] == reg.id && !i[:registered] }
+        end
+        return true
+      end
+
+      original_request['params'].unregisterations.each do |reg|
+        PuppetLanguageServer.log_message(:info, "Succesfully dynamically unregistered the #{reg.method__lsp} method")
+
+        # Remove registrations
+        @registrations[reg.method__lsp] = [] if @registrations[reg.method__lsp].nil?
+        @registrations[reg.method__lsp].delete_if { |i| i[:id] == reg.id }
       end
 
       true
     end
 
     private
 
+    def new_request_id
+      SecureRandom.uuid
+    end
+
     def safe_hash_traverse(hash, *names)
       return nil if names.empty?
       item = nil

lib/puppet-languageserver/message_router.rb

@@ -34,7 +34,7 @@ class MessageRouter < BaseMessageRouter
     def initialize(options = {})
       super
       @server_options = options.nil? ? {} : options
-      @client = LanguageClient.new
+      @client = LanguageClient.new(self)
     end
 
     def documents
@@ -254,7 +254,7 @@ def receive_notification(method, params)
       when 'initialized'
         PuppetLanguageServer.log_message(:info, 'Client has received initialization')
         if client.client_capability('workspace', 'didChangeConfiguration', 'dynamicRegistration') == true
-          client.register_capability(self, 'workspace/didChangeConfiguration')
+          client.register_capability('workspace/didChangeConfiguration')
         else
           PuppetLanguageServer.log_message(:debug, 'Client does not support didChangeConfiguration dynamic registration. Using push method for configuration change detection.')
         end
@@ -300,7 +300,7 @@ def receive_notification(method, params)
         if params.key?('settings') && params['settings'].nil?
           # This is a notification from a dynamic registration.  Need to send a workspace/configuration
           # request to get the actual configuration
-          client.send_configuration_request(self)
+          client.send_configuration_request
         else
           client.parse_lsp_configuration_settings!(params['settings'])
         end
@@ -314,15 +314,17 @@ def receive_notification(method, params)
     end
 
     def receive_response(response, original_request)
-      unless response.key?('result')
+      unless receive_response_succesful?(response) # rubocop:disable Style/IfUnlessModifier Line is too long otherwise
         PuppetLanguageServer.log_message(:error, "Response for method '#{original_request['method']}' with id '#{original_request['id']}' failed with #{response['error']}")
-        return
       end
-
+      # Error responses still need to be processed so process messages even if it failed
       case original_request['method']
       when 'client/registerCapability'
-        client.parse_register_capability_response!(self, response, original_request)
+        client.parse_register_capability_response!(response, original_request)
+      when 'client/unregisterCapability'
+        client.parse_unregister_capability_response!(response, original_request)
       when 'workspace/configuration'
+        return unless receive_response_succesful?(response)
         client.parse_lsp_configuration_settings!(response['result'])
       else
         super
@@ -331,6 +333,12 @@ def receive_response(response, original_request)
       PuppetLanguageServer::CrashDump.write_crash_file(e, nil, 'response' => response, 'original_request' => original_request)
       raise
     end
+
+    private
+
+    def receive_response_succesful?(response)
+      response.key?('result')
+    end
   end
 
   class DisabledMessageRouter < BaseMessageRouter