Add support for delegating completion requests for ERB files (Shopify…

…#2551)

* Add support for delegating completion requests for ERB files

* Rename parse to parse!

* Add more delegation documentation in the code

lib/ruby_lsp/base_server.rb

@@ -31,6 +31,7 @@ def initialize(test_mode: false)
         Thread,
       )
 
+      @global_state = T.let(GlobalState.new, GlobalState)
       Thread.main.priority = 1
     end
 
@@ -52,7 +53,26 @@ def start
               message[:params][:textDocument][:uri] = parsed_uri
 
               # We don't want to try to parse documents on text synchronization notifications
-              @store.get(parsed_uri).parse unless method.start_with?("textDocument/did")
+              unless method.start_with?("textDocument/did")
+                document = @store.get(parsed_uri)
+
+                # If the client supports request delegation and we're working with an ERB document and there was
+                # something to parse, then we have to maintain the client updated about the virtual state of the host
+                # language source
+                if document.parse! && @global_state.supports_request_delegation && document.is_a?(ERBDocument)
+                  send_message(
+                    Notification.new(
+                      method: "delegate/textDocument/virtualState",
+                      params: {
+                        textDocument: {
+                          uri: uri,
+                          text: document.host_language_source,
+                        },
+                      },
+                    ),
+                  )
+                end
+              end
             rescue Store::NonExistingDocumentError
               # If we receive a request for a file that no longer exists, we don't want to fail
             end

lib/ruby_lsp/document.rb

@@ -51,7 +51,8 @@ def initialize(source:, version:, uri:, encoding: Encoding::UTF_8)
       @version = T.let(version, Integer)
       @uri = T.let(uri, URI::Generic)
       @needs_parsing = T.let(true, T::Boolean)
-      @parse_result = T.let(parse, ParseResultType)
+      @parse_result = T.let(T.unsafe(nil), ParseResultType)
+      parse!
     end
 
     sig { params(other: Document[T.untyped]).returns(T::Boolean) }
@@ -106,8 +107,9 @@ def push_edits(edits, version:)
       @cache.clear
     end
 
-    sig { abstract.returns(ParseResultType) }
-    def parse; end
+    # Returns `true` if the document was parsed and `false` if nothing needed parsing
+    sig { abstract.returns(T::Boolean) }
+    def parse!; end
 
     sig { abstract.returns(T::Boolean) }
     def syntax_error?; end

lib/ruby_lsp/erb_document.rb

@@ -6,17 +6,30 @@ class ERBDocument < Document
     extend T::Sig
     extend T::Generic
 
+    sig { returns(String) }
+    attr_reader :host_language_source
+
     ParseResultType = type_member { { fixed: Prism::ParseResult } }
 
-    sig { override.returns(ParseResultType) }
-    def parse
-      return @parse_result unless @needs_parsing
+    sig { params(source: String, version: Integer, uri: URI::Generic, encoding: Encoding).void }
+    def initialize(source:, version:, uri:, encoding: Encoding::UTF_8)
+      # This has to be initialized before calling super because we call `parse` in the parent constructor, which
+      # overrides this with the proper virtual host language source
+      @host_language_source = T.let("", String)
+      super
+    end
+
+    sig { override.returns(T::Boolean) }
+    def parse!
+      return false unless @needs_parsing
 
       @needs_parsing = false
       scanner = ERBScanner.new(@source)
       scanner.scan
+      @host_language_source = scanner.host_language
       # assigning empty scopes to turn Prism into eval mode
       @parse_result = Prism.parse(scanner.ruby, scopes: [[]])
+      true
     end
 
     sig { override.returns(T::Boolean) }
@@ -39,16 +52,22 @@ def locate_node(position, node_types: [])
       RubyDocument.locate(@parse_result.value, create_scanner.find_char_position(position), node_types: node_types)
     end
 
+    sig { params(char_position: Integer).returns(T.nilable(T::Boolean)) }
+    def inside_host_language?(char_position)
+      char = @host_language_source[char_position]
+      char && char != " "
+    end
+
     class ERBScanner
       extend T::Sig
 
       sig { returns(String) }
-      attr_reader :ruby, :html
+      attr_reader :ruby, :host_language
 
       sig { params(source: String).void }
       def initialize(source)
         @source = source
-        @html = T.let(+"", String)
+        @host_language = T.let(+"", String)
         @ruby = T.let(+"", String)
         @current_pos = T.let(0, Integer)
         @inside_ruby = T.let(false, T::Boolean)
@@ -104,16 +123,16 @@ def scan_char
           end
         when "\r"
           @ruby << char
-          @html << char
+          @host_language << char
 
           if next_char == "\n"
             @ruby << next_char
-            @html << next_char
+            @host_language << next_char
             @current_pos += 1
           end
         when "\n"
           @ruby << char
-          @html << char
+          @host_language << char
         else
           push_char(T.must(char))
         end
@@ -123,10 +142,10 @@ def scan_char
       def push_char(char)
         if @inside_ruby
           @ruby << char
-          @html << " " * char.length
+          @host_language << " " * char.length
         else
           @ruby << " " * char.length
-          @html << char
+          @host_language << char
         end
       end
 

lib/ruby_lsp/global_state.rb

@@ -21,7 +21,7 @@ class GlobalState
     attr_reader :encoding
 
     sig { returns(T::Boolean) }
-    attr_reader :supports_watching_files, :experimental_features
+    attr_reader :supports_watching_files, :experimental_features, :supports_request_delegation
 
     sig { returns(TypeInferrer) }
     attr_reader :type_inferrer
@@ -41,6 +41,7 @@ def initialize
       @experimental_features = T.let(false, T::Boolean)
       @type_inferrer = T.let(TypeInferrer.new(@index, @experimental_features), TypeInferrer)
       @addon_settings = T.let({}, T::Hash[String, T.untyped])
+      @supports_request_delegation = T.let(false, T::Boolean)
     end
 
     sig { params(addon_name: String).returns(T.nilable(T::Hash[Symbol, T.untyped])) }
@@ -131,6 +132,7 @@ def apply_options(options)
         @addon_settings.merge!(addon_settings)
       end
 
+      @supports_request_delegation = options.dig(:capabilities, :experimental, :requestDelegation) || false
       notifications
     end
 

lib/ruby_lsp/rbs_document.rb

@@ -14,18 +14,19 @@ def initialize(source:, version:, uri:, encoding: Encoding::UTF_8)
       super
     end
 
-    sig { override.returns(ParseResultType) }
-    def parse
-      return @parse_result unless @needs_parsing
+    sig { override.returns(T::Boolean) }
+    def parse!
+      return false unless @needs_parsing
 
       @needs_parsing = false
 
       _, _, declarations = RBS::Parser.parse_signature(@source)
       @syntax_error = false
       @parse_result = declarations
+      true
     rescue RBS::ParsingError
       @syntax_error = true
-      @parse_result
+      true
     end
 
     sig { override.returns(T::Boolean) }

lib/ruby_lsp/requests/completion.rb

@@ -60,6 +60,8 @@ def initialize(document, global_state, params, sorbet_level, dispatcher)
         # Completion always receives the position immediately after the character that was just typed. Here we adjust it
         # back by 1, so that we find the right node
         char_position = document.create_scanner.find_char_position(params[:position]) - 1
+        delegate_request_if_needed!(global_state, document, char_position)
+
         node_context = RubyDocument.locate(
           document.parse_result.value,
           char_position,

lib/ruby_lsp/requests/request.rb

@@ -17,6 +17,23 @@ def perform; end
 
       private
 
+      # Signals to the client that the request should be delegated to the language server server for the host language
+      # in ERB files
+      sig do
+        params(
+          global_state: GlobalState,
+          document: Document[T.untyped],
+          char_position: Integer,
+        ).void
+      end
+      def delegate_request_if_needed!(global_state, document, char_position)
+        if global_state.supports_request_delegation &&
+            document.is_a?(ERBDocument) &&
+            document.inside_host_language?(char_position)
+          raise DelegateRequestError
+        end
+      end
+
       # Checks if a location covers a position
       sig { params(location: Prism::Location, position: T.untyped).returns(T::Boolean) }
       def cover?(location, position)

lib/ruby_lsp/ruby_document.rb

@@ -121,12 +121,13 @@ def locate(node, char_position, node_types: [])
       end
     end
 
-    sig { override.returns(ParseResultType) }
-    def parse
-      return @parse_result unless @needs_parsing
+    sig { override.returns(T::Boolean) }
+    def parse!
+      return false unless @needs_parsing
 
       @needs_parsing = false
       @parse_result = Prism.parse(@source)
+      true
     end
 
     sig { override.returns(T::Boolean) }

lib/ruby_lsp/server.rb

@@ -9,12 +9,6 @@ class Server < BaseServer
     sig { returns(GlobalState) }
     attr_reader :global_state
 
-    sig { params(test_mode: T::Boolean).void }
-    def initialize(test_mode: false)
-      super
-      @global_state = T.let(GlobalState.new, GlobalState)
-    end
-
     sig { override.params(message: T::Hash[Symbol, T.untyped]).void }
     def process_message(message)
       case message[:method]
@@ -98,6 +92,8 @@ def process_message(message)
       when "$/cancelRequest"
         @mutex.synchronize { @cancelled_requests << message[:params][:id] }
       end
+    rescue DelegateRequestError
+      send_message(Error.new(id: message[:id], code: DelegateRequestError::CODE, message: "DELEGATE_REQUEST"))
     rescue StandardError, LoadError => e
       # If an error occurred in a request, we have to return an error response or else the editor will hang
       if message[:id]
@@ -748,6 +744,17 @@ def text_document_completion(message)
 
     sig { params(message: T::Hash[Symbol, T.untyped]).void }
     def text_document_completion_item_resolve(message)
+      # When responding to a delegated completion request, it means we're handling a completion item that isn't related
+      # to Ruby (probably related to an ERB host language like HTML). We need to return the original completion item
+      # back to the editor so that it's displayed correctly
+      if message.dig(:params, :data, :delegateCompletion)
+        send_message(Result.new(
+          id: message[:id],
+          response: message[:params],
+        ))
+        return
+      end
+
       send_message(Result.new(
         id: message[:id],
         response: Requests::CompletionResolve.new(@global_state, message[:params]).perform,

lib/ruby_lsp/utils.rb

@@ -27,6 +27,16 @@ module RubyLsp
   )
   GUESSED_TYPES_URL = "https://github.com/Shopify/ruby-lsp/blob/main/DESIGN_AND_ROADMAP.md#guessed-types"
 
+  # Request delegation for embedded languages is not yet standardized into the language server specification. Here we
+  # use this custom error class as a way to return a signal to the client that the request should be delegated to the
+  # language server for the host language. The support for delegation is custom built on the client side, so each editor
+  # needs to implement their own until this becomes a part of the spec
+  class DelegateRequestError < StandardError
+    # A custom error code that clients can use to handle delegate requests. This is past the range of error codes listed
+    # by the specification to avoid conflicting with other error types
+    CODE = -32900
+  end
+
   # A notification to be sent to the client
   class Message
     extend T::Sig

test/erb_document_test.rb

@@ -14,7 +14,7 @@ def test_erb_file_is_properly_parsed
       </ul>
     ERB
 
-    document.parse
+    document.parse!
 
     refute_predicate(document, :syntax_error?)
     assert_equal(
@@ -35,7 +35,7 @@ def test_erb_file_parses_in_eval_context
       </html>
     ERB
 
-    document.parse
+    document.parse!
 
     refute_predicate(document, :syntax_error?)
     assert_equal(
@@ -46,7 +46,7 @@ def test_erb_file_parses_in_eval_context
 
   def test_erb_document_handles_windows_newlines
     document = RubyLsp::ERBDocument.new(source: "<%=\r\nbar %>", version: 1, uri: URI("file:///foo.erb"))
-    document.parse
+    document.parse!
 
     refute_predicate(document, :syntax_error?)
     assert_equal("   \r\nbar   ", document.parse_result.source.source)
@@ -62,7 +62,7 @@ def test_erb_syntax_error_doesnt_cause_crash
       "<%= foo %\n<%= bar %>",
     ].each do |source|
       document = RubyLsp::ERBDocument.new(source: source, version: 1, uri: URI("file:///foo.erb"))
-      document.parse
+      document.parse!
     end
   end
 
@@ -106,4 +106,20 @@ def test_cache_set_and_get
     assert_equal(value, document.cache_set("textDocument/semanticHighlighting", value))
     assert_equal(value, document.cache_get("textDocument/semanticHighlighting"))
   end
+
+  def test_keeps_track_of_virtual_host_language_source
+    document = RubyLsp::ERBDocument.new(source: +<<~ERB, version: 1, uri: URI("file:///foo.erb"))
+      <ul>
+        <li><%= foo %><li>
+        <li><%= end %><li>
+      </ul>
+    ERB
+
+    assert_equal(<<~HTML, document.host_language_source)
+      <ul>
+        <li>          <li>
+        <li>          <li>
+      </ul>
+    HTML
+  end
 end

test/rbs_document_test.rb

@@ -29,7 +29,7 @@ def bar: () - void
       [{ range: { start: { line: 1, character: 15 }, end: { line: 1, character: 15 } }, text: ">" }],
       version: 2,
     )
-    document.parse
+    document.parse!
     refute_predicate(document, :syntax_error?)
   end
 end