Add support for delegating completion requests for ERB files

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)
+                document.parse
+
+                # If the client supports request delegation and we're working with an ERB document, then we have to
+                # maintain the client updated about the virtual state of the host language source
+                if @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/erb_document.rb

@@ -6,15 +6,27 @@ class ERBDocument < Document
     extend T::Sig
     extend T::Generic
 
+    sig { returns(String) }
+    attr_reader :host_language_source
+
     ParseResultType = type_member { { fixed: Prism::ParseResult } }
 
+    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(ParseResultType) }
     def parse
       return @parse_result 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: [[]])
     end
@@ -39,16 +51,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 +122,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 +141,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/requests/completion.rb

@@ -60,6 +60,11 @@ 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
+
+        if document.is_a?(ERBDocument) && document.inside_host_language?(char_position)
+          raise DelegateRequestError
+        end
+
         node_context = RubyDocument.locate(
           document.parse_result.value,
           char_position,

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: -32900, 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,9 @@ module RubyLsp
   )
   GUESSED_TYPES_URL = "https://github.com/Shopify/ruby-lsp/blob/main/DESIGN_AND_ROADMAP.md#guessed-types"
 
+  class DelegateRequestError < StandardError
+  end
+
   # A notification to be sent to the client
   class Message
     extend T::Sig

test/erb_document_test.rb

@@ -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

vscode/src/client.ts

@@ -22,6 +22,11 @@ import {
   CloseAction,
   State,
   DocumentFilter,
+  CompletionList,
+  StaticFeature,
+  ClientCapabilities,
+  FeatureState,
+  ServerCapabilities,
 } from "vscode-languageclient/node";
 
 import {
@@ -224,6 +229,25 @@ class ClientErrorHandler implements ErrorHandler {
   }
 }
 
+class ExperimentalCapabilities implements StaticFeature {
+  fillClientCapabilities(capabilities: ClientCapabilities): void {
+    capabilities.experimental = {
+      requestDelegation: true,
+    };
+  }
+
+  initialize(
+    _capabilities: ServerCapabilities,
+    _documentSelector: DocumentSelector | undefined,
+  ): void {}
+
+  getState(): FeatureState {
+    return { kind: "static" };
+  }
+
+  clear(): void {}
+}
+
 export default class Client extends LanguageClient implements ClientInterface {
   public readonly ruby: Ruby;
   public serverVersion?: string;
@@ -233,6 +257,7 @@ export default class Client extends LanguageClient implements ClientInterface {
   private readonly createTestItems: (response: CodeLens[]) => void;
   private readonly baseFolder;
   private readonly workspaceOutputChannel: WorkspaceChannel;
+  private readonly virtualDocuments = new Map<string, string>();
 
   #context: vscode.ExtensionContext;
   #formatter: string;
@@ -244,6 +269,7 @@ export default class Client extends LanguageClient implements ClientInterface {
     createTestItems: (response: CodeLens[]) => void,
     workspaceFolder: vscode.WorkspaceFolder,
     outputChannel: WorkspaceChannel,
+    virtualDocuments: Map<string, string>,
     isMainWorkspace = false,
     debugMode?: boolean,
   ) {
@@ -260,7 +286,9 @@ export default class Client extends LanguageClient implements ClientInterface {
       debugMode,
     );
 
+    this.registerFeature(new ExperimentalCapabilities());
     this.workspaceOutputChannel = outputChannel;
+    this.virtualDocuments = virtualDocuments;
 
     // Middleware are part of client options, but because they must reference `this`, we cannot make it a part of the
     // `super` call (TypeScript does not allow accessing `this` before invoking `super`)
@@ -273,6 +301,13 @@ export default class Client extends LanguageClient implements ClientInterface {
     this.#context = context;
     this.ruby = ruby;
     this.#formatter = "";
+
+    this.onNotification("delegate/textDocument/virtualState", (params) => {
+      this.virtualDocuments.set(
+        params.textDocument.uri,
+        params.textDocument.text,
+      );
+    });
   }
 
   async afterStart() {
@@ -317,7 +352,7 @@ export default class Client extends LanguageClient implements ClientInterface {
 
   private async benchmarkMiddleware<T>(
     type: string | MessageSignature,
-    _params: any,
+    params: any,
     runRequest: () => Promise<T>,
   ): Promise<T> {
     if (this.state !== State.Running) {
@@ -341,6 +376,12 @@ export default class Client extends LanguageClient implements ClientInterface {
       this.logResponseTime(bench.duration, request);
       return result;
     } catch (error: any) {
+      // We use a special error code to indicate delegated requests. This is not actually an error, it's a signal that
+      // the client needs to invoke the appropriate language service for this request
+      if (error.code === -32900) {
+        return this.executeDelegateRequest(type, params);
+      }
+
       if (error.data) {
         if (
           this.baseFolder === "ruby-lsp" ||
@@ -379,6 +420,58 @@ export default class Client extends LanguageClient implements ClientInterface {
     }
   }
 
+  private async executeDelegateRequest(
+    type: string | MessageSignature,
+    params: any,
+  ): Promise<any> {
+    const request = typeof type === "string" ? type : type.method;
+    const originalUri = params.textDocument.uri;
+
+    // Delegating requests only makes sense for text document requests, where a URI is available
+    if (!originalUri) {
+      return null;
+    }
+
+    const hostLanguage = /\.([^.]+)\.erb$/.exec(originalUri)?.[1] || "html";
+    const vdocUriString = `embedded-content://${hostLanguage}/${encodeURIComponent(
+      originalUri,
+    )}.${hostLanguage}`;
+
+    if (request === "textDocument/completion") {
+      return vscode.commands
+        .executeCommand<CompletionList>(
+          "vscode.executeCompletionItemProvider",
+          vscode.Uri.parse(vdocUriString),
+          params.position,
+          params.context.triggerCharacter,
+        )
+        .then((response) => {
+          // We need to tell the server that the completion item is being delegated, so that when it receives the
+          // `completionItem/resolve`, we can delegate that too
+          response.items.forEach((item) => {
+            // For whatever reason, HTML completion items don't include the `kind` and that causes a failure in the
+            // editor. It might be a mistake in the delegation
+            if (
+              item.documentation &&
+              typeof item.documentation !== "string" &&
+              "value" in item.documentation
+            ) {
+              item.documentation.kind = "markdown";
+            }
+
+            item.data = { ...item.data, delegateCompletion: true };
+          });
+
+          return response;
+        });
+    } else {
+      this.workspaceOutputChannel.warn(
+        `Attempted to delegate unsupported request ${request}`,
+      );
+      return null;
+    }
+  }
+
   // Register the middleware in the client options
   private registerMiddleware() {
     this.clientOptions.middleware = {