ruby-lsp 0.17.13 → 0.17.15

data/lib/ruby_lsp/requests/document_highlight.rb

@@ -29,7 +29,7 @@ module RubyLsp
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           position: T::Hash[Symbol, T.untyped],
           dispatcher: Prism::Dispatcher,
         ).void

data/lib/ruby_lsp/requests/formatting.rb

@@ -40,7 +40,7 @@ module RubyLsp
         end
       end
 
-      sig { params(global_state: GlobalState, document: Document).void }
+      sig { params(global_state: GlobalState, document: RubyDocument).void }
       def initialize(global_state, document)
         super()
         @document = document

data/lib/ruby_lsp/requests/hover.rb

@@ -32,7 +32,7 @@ module RubyLsp
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           global_state: GlobalState,
           position: T::Hash[Symbol, T.untyped],
           dispatcher: Prism::Dispatcher,

data/lib/ruby_lsp/requests/inlay_hints.rb

@@ -52,7 +52,7 @@ module RubyLsp
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           range: T::Hash[Symbol, T.untyped],
           hints_configuration: RequestConfig,
           dispatcher: Prism::Dispatcher,

data/lib/ruby_lsp/requests/on_type_formatting.rb

@@ -42,7 +42,7 @@ module RubyLsp
 
       sig do
         params(
-          document: Document,
+          document: RubyDocument,
           position: T::Hash[Symbol, T.untyped],
           trigger_character: String,
           client_name: String,

data/lib/ruby_lsp/requests/prepare_type_hierarchy.rb

@@ -35,7 +35,7 @@ module RubyLsp
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           index: RubyIndexer::Index,
           position: T::Hash[Symbol, T.untyped],
         ).void

data/lib/ruby_lsp/requests/selection_ranges.rb

@@ -23,7 +23,8 @@ module RubyLsp
     class SelectionRanges < Request
       extend T::Sig
       include Support::Common
-      sig { params(document: Document).void }
+
+      sig { params(document: T.any(RubyDocument, ERBDocument)).void }
       def initialize(document)
         super()
         @document = document

data/lib/ruby_lsp/requests/show_syntax_tree.rb

@@ -20,7 +20,7 @@ module RubyLsp
     class ShowSyntaxTree < Request
       extend T::Sig
 
-      sig { params(document: Document, range: T.nilable(T::Hash[Symbol, T.untyped])).void }
+      sig { params(document: RubyDocument, range: T.nilable(T::Hash[Symbol, T.untyped])).void }
       def initialize(document, range)
         super()
         @document = document

data/lib/ruby_lsp/requests/signature_help.rb

@@ -41,7 +41,7 @@ module RubyLsp
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           global_state: GlobalState,
           position: T::Hash[Symbol, T.untyped],
           context: T.nilable(T::Hash[Symbol, T.untyped]),

data/lib/ruby_lsp/requests/support/formatter.rb

@@ -10,13 +10,13 @@ module RubyLsp
 
         interface!
 
-        sig { abstract.params(uri: URI::Generic, document: Document).returns(T.nilable(String)) }
+        sig { abstract.params(uri: URI::Generic, document: RubyDocument).returns(T.nilable(String)) }
         def run_formatting(uri, document); end
 
         sig do
           abstract.params(
             uri: URI::Generic,
-            document: Document,
+            document: RubyDocument,
           ).returns(T.nilable(T::Array[Interface::Diagnostic]))
         end
         def run_diagnostic(uri, document); end

data/lib/ruby_lsp/requests/support/rubocop_diagnostic.rb

@@ -31,7 +31,7 @@ module RubyLsp
 
         # TODO: avoid passing document once we have alternative ways to get at
         # encoding and file source
-        sig { params(document: Document, offense: RuboCop::Cop::Offense, uri: URI::Generic).void }
+        sig { params(document: RubyDocument, offense: RuboCop::Cop::Offense, uri: URI::Generic).void }
         def initialize(document, offense, uri)
           @document = document
           @offense = offense

data/lib/ruby_lsp/requests/support/rubocop_formatter.rb

@@ -17,7 +17,7 @@ module RubyLsp
           @format_runner = T.let(RuboCopRunner.new("-a"), RuboCopRunner)
         end
 
-        sig { override.params(uri: URI::Generic, document: Document).returns(T.nilable(String)) }
+        sig { override.params(uri: URI::Generic, document: RubyDocument).returns(T.nilable(String)) }
         def run_formatting(uri, document)
           filename = T.must(uri.to_standardized_path || uri.opaque)
 
@@ -29,7 +29,7 @@ module RubyLsp
         sig do
           override.params(
             uri: URI::Generic,
-            document: Document,
+            document: RubyDocument,
           ).returns(T.nilable(T::Array[Interface::Diagnostic]))
         end
         def run_diagnostic(uri, document)

data/lib/ruby_lsp/requests/support/syntax_tree_formatter.rb

@@ -29,7 +29,7 @@ module RubyLsp
             )
         end
 
-        sig { override.params(uri: URI::Generic, document: Document).returns(T.nilable(String)) }
+        sig { override.params(uri: URI::Generic, document: RubyDocument).returns(T.nilable(String)) }
         def run_formatting(uri, document)
           path = uri.to_standardized_path
           return if path && @options.ignore_files.any? { |pattern| File.fnmatch?("*/#{pattern}", path) }
@@ -40,7 +40,7 @@ module RubyLsp
         sig do
           override.params(
             uri: URI::Generic,
-            document: Document,
+            document: RubyDocument,
           ).returns(T.nilable(T::Array[Interface::Diagnostic]))
         end
         def run_diagnostic(uri, document)

data/lib/ruby_lsp/ruby_document.rb

@@ -3,6 +3,11 @@
 
 module RubyLsp
   class RubyDocument < Document
+    extend T::Sig
+    extend T::Generic
+
+    ParseResultType = type_member { { fixed: Prism::ParseResult } }
+
     class SorbetLevel < T::Enum
       enums do
         None = new("none")
@@ -13,7 +18,110 @@ module RubyLsp
       end
     end
 
-    sig { override.returns(Prism::ParseResult) }
+    class << self
+      extend T::Sig
+
+      sig do
+        params(
+          node: Prism::Node,
+          char_position: Integer,
+          node_types: T::Array[T.class_of(Prism::Node)],
+        ).returns(NodeContext)
+      end
+      def locate(node, char_position, node_types: [])
+        queue = T.let(node.child_nodes.compact, T::Array[T.nilable(Prism::Node)])
+        closest = node
+        parent = T.let(nil, T.nilable(Prism::Node))
+        nesting_nodes = T.let(
+          [],
+          T::Array[T.any(
+            Prism::ClassNode,
+            Prism::ModuleNode,
+            Prism::SingletonClassNode,
+            Prism::DefNode,
+            Prism::BlockNode,
+            Prism::LambdaNode,
+            Prism::ProgramNode,
+          )],
+        )
+
+        nesting_nodes << node if node.is_a?(Prism::ProgramNode)
+        call_node = T.let(nil, T.nilable(Prism::CallNode))
+
+        until queue.empty?
+          candidate = queue.shift
+
+          # Skip nil child nodes
+          next if candidate.nil?
+
+          # Add the next child_nodes to the queue to be processed. The order here is important! We want to move in the
+          # same order as the visiting mechanism, which means searching the child nodes before moving on to the next
+          # sibling
+          T.unsafe(queue).unshift(*candidate.child_nodes)
+
+          # Skip if the current node doesn't cover the desired position
+          loc = candidate.location
+          next unless (loc.start_offset...loc.end_offset).cover?(char_position)
+
+          # If the node's start character is already past the position, then we should've found the closest node
+          # already
+          break if char_position < loc.start_offset
+
+          # If the candidate starts after the end of the previous nesting level, then we've exited that nesting level
+          # and need to pop the stack
+          previous_level = nesting_nodes.last
+          nesting_nodes.pop if previous_level && loc.start_offset > previous_level.location.end_offset
+
+          # Keep track of the nesting where we found the target. This is used to determine the fully qualified name of
+          # the target when it is a constant
+          case candidate
+          when Prism::ClassNode, Prism::ModuleNode, Prism::SingletonClassNode, Prism::DefNode, Prism::BlockNode,
+            Prism::LambdaNode
+            nesting_nodes << candidate
+          end
+
+          if candidate.is_a?(Prism::CallNode)
+            arg_loc = candidate.arguments&.location
+            blk_loc = candidate.block&.location
+            if (arg_loc && (arg_loc.start_offset...arg_loc.end_offset).cover?(char_position)) ||
+                (blk_loc && (blk_loc.start_offset...blk_loc.end_offset).cover?(char_position))
+              call_node = candidate
+            end
+          end
+
+          # If there are node types to filter by, and the current node is not one of those types, then skip it
+          next if node_types.any? && node_types.none? { |type| candidate.class == type }
+
+          # If the current node is narrower than or equal to the previous closest node, then it is more precise
+          closest_loc = closest.location
+          if loc.end_offset - loc.start_offset <= closest_loc.end_offset - closest_loc.start_offset
+            parent = closest
+            closest = candidate
+          end
+        end
+
+        # When targeting the constant part of a class/module definition, we do not want the nesting to be duplicated.
+        # That is, when targeting Bar in the following example:
+        #
+        # ```ruby
+        #   class Foo::Bar; end
+        # ```
+        # The correct target is `Foo::Bar` with an empty nesting. `Foo::Bar` should not appear in the nesting stack,
+        # even though the class/module node does indeed enclose the target, because it would lead to incorrect behavior
+        if closest.is_a?(Prism::ConstantReadNode) || closest.is_a?(Prism::ConstantPathNode)
+          last_level = nesting_nodes.last
+
+          if (last_level.is_a?(Prism::ModuleNode) || last_level.is_a?(Prism::ClassNode)) &&
+              last_level.constant_path == closest
+            nesting_nodes.pop
+          end
+        end
+
+        NodeContext.new(closest, parent, nesting_nodes, call_node)
+      end
+    end
+
+    sig { override.returns(ParseResultType) }
     def parse
       return @parse_result unless @needs_parsing
 
@@ -84,5 +192,15 @@ module RubyLsp
         end
       end
     end
+
+    sig do
+      params(
+        position: T::Hash[Symbol, T.untyped],
+        node_types: T::Array[T.class_of(Prism::Node)],
+      ).returns(NodeContext)
+    end
+    def locate_node(position, node_types: [])
+      RubyDocument.locate(@parse_result.value, create_scanner.find_char_position(position), node_types: node_types)
+    end
   end
 end

data/lib/ruby_lsp/server.rb

@@ -298,9 +298,12 @@ module RubyLsp
         language_id = case text_document[:languageId]
         when "erb", "eruby"
           Document::LanguageId::ERB
+        when "rbs"
+          Document::LanguageId::RBS
         else
           Document::LanguageId::Ruby
         end
+
         @store.set(
           uri: text_document[:uri],
           source: text_document[:text],
@@ -341,7 +344,12 @@ module RubyLsp
     def text_document_selection_range(message)
       uri = message.dig(:params, :textDocument, :uri)
       ranges = @store.cache_fetch(uri, "textDocument/selectionRange") do |document|
-        Requests::SelectionRanges.new(document).perform
+        case document
+        when RubyDocument, ERBDocument
+          Requests::SelectionRanges.new(document).perform
+        else
+          []
+        end
       end
 
       # Per the selection range request spec (https://microsoft.github.io/language-server-protocol/specification#textDocument_selectionRange),
@@ -363,6 +371,11 @@ module RubyLsp
       uri = URI(message.dig(:params, :textDocument, :uri))
       document = @store.get(uri)
 
+      unless document.is_a?(RubyDocument) || document.is_a?(ERBDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
       # If the response has already been cached by another request, return it
       cached_response = document.cache_get(message[:method])
       if cached_response
@@ -407,6 +420,12 @@ module RubyLsp
       range = params[:range]
       uri = params.dig(:textDocument, :uri)
       document = @store.get(uri)
+
+      unless document.is_a?(RubyDocument) || document.is_a?(ERBDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
       start_line = range.dig(:start, :line)
       end_line = range.dig(:end, :line)
 
@@ -436,6 +455,10 @@ module RubyLsp
       end
 
       document = @store.get(uri)
+      unless document.is_a?(RubyDocument)
+        send_empty_response(message[:id])
+        return
+      end
 
       response = Requests::Formatting.new(@global_state, document).perform
       send_message(Result.new(id: message[:id], response: response))
@@ -452,6 +475,12 @@ module RubyLsp
       params = message[:params]
       dispatcher = Prism::Dispatcher.new
       document = @store.get(params.dig(:textDocument, :uri))
+
+      unless document.is_a?(RubyDocument) || document.is_a?(ERBDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
       request = Requests::DocumentHighlight.new(document, params[:position], dispatcher)
       dispatcher.dispatch(document.parse_result.value)
       send_message(Result.new(id: message[:id], response: request.perform))
@@ -462,6 +491,11 @@ module RubyLsp
       params = message[:params]
       document = @store.get(params.dig(:textDocument, :uri))
 
+      unless document.is_a?(RubyDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
       send_message(
         Result.new(
           id: message[:id],
@@ -481,6 +515,11 @@ module RubyLsp
       dispatcher = Prism::Dispatcher.new
       document = @store.get(params.dig(:textDocument, :uri))
 
+      unless document.is_a?(RubyDocument) || document.is_a?(ERBDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
       send_message(
         Result.new(
           id: message[:id],
@@ -495,7 +534,7 @@ module RubyLsp
       )
     end
 
-    sig { params(document: Document).returns(RubyDocument::SorbetLevel) }
+    sig { params(document: Document[T.untyped]).returns(RubyDocument::SorbetLevel) }
     def sorbet_level(document)
       return RubyDocument::SorbetLevel::Ignore unless @global_state.has_type_checker
       return RubyDocument::SorbetLevel::Ignore unless document.is_a?(RubyDocument)
@@ -509,6 +548,12 @@ module RubyLsp
       hints_configurations = T.must(@store.features_configuration.dig(:inlayHint))
       dispatcher = Prism::Dispatcher.new
       document = @store.get(params.dig(:textDocument, :uri))
+
+      unless document.is_a?(RubyDocument) || document.is_a?(ERBDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
       request = Requests::InlayHints.new(document, params[:range], hints_configurations, dispatcher)
       dispatcher.visit(document.parse_result.value)
       send_message(Result.new(id: message[:id], response: request.perform))
@@ -519,6 +564,11 @@ module RubyLsp
       params = message[:params]
       document = @store.get(params.dig(:textDocument, :uri))
 
+      unless document.is_a?(RubyDocument) || document.is_a?(ERBDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
       send_message(
         Result.new(
           id: message[:id],
@@ -581,7 +631,10 @@ module RubyLsp
       document = @store.get(uri)
 
       response = document.cache_fetch("textDocument/diagnostic") do |document|
-        Requests::Diagnostics.new(@global_state, document).perform
+        case document
+        when RubyDocument
+          Requests::Diagnostics.new(@global_state, document).perform
+        end
       end
 
       send_message(
@@ -604,6 +657,11 @@ module RubyLsp
       dispatcher = Prism::Dispatcher.new
       document = @store.get(params.dig(:textDocument, :uri))
 
+      unless document.is_a?(RubyDocument) || document.is_a?(ERBDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
       send_message(
         Result.new(
           id: message[:id],
@@ -632,6 +690,11 @@ module RubyLsp
       dispatcher = Prism::Dispatcher.new
       document = @store.get(params.dig(:textDocument, :uri))
 
+      unless document.is_a?(RubyDocument) || document.is_a?(ERBDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
       send_message(
         Result.new(
           id: message[:id],
@@ -653,6 +716,11 @@ module RubyLsp
       dispatcher = Prism::Dispatcher.new
       document = @store.get(params.dig(:textDocument, :uri))
 
+      unless document.is_a?(RubyDocument) || document.is_a?(ERBDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
       send_message(
         Result.new(
           id: message[:id],
@@ -710,9 +778,16 @@ module RubyLsp
     sig { params(message: T::Hash[Symbol, T.untyped]).void }
     def text_document_show_syntax_tree(message)
       params = message[:params]
+      document = @store.get(params.dig(:textDocument, :uri))
+
+      unless document.is_a?(RubyDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
       response = {
         ast: Requests::ShowSyntaxTree.new(
-          @store.get(params.dig(:textDocument, :uri)),
+          document,
           params[:range],
         ).perform,
       }
@@ -722,11 +797,19 @@ module RubyLsp
     sig { params(message: T::Hash[Symbol, T.untyped]).void }
     def text_document_prepare_type_hierarchy(message)
       params = message[:params]
+      document = @store.get(params.dig(:textDocument, :uri))
+
+      unless document.is_a?(RubyDocument) || document.is_a?(ERBDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
       response = Requests::PrepareTypeHierarchy.new(
-        @store.get(params.dig(:textDocument, :uri)),
+        document,
         @global_state.index,
         params[:position],
       ).perform
+
       send_message(Result.new(id: message[:id], response: response))
     end
 
@@ -793,6 +876,11 @@ module RubyLsp
           send_message(Notification.window_show_error("Error while indexing: #{error.message}"))
         end
 
+        # Indexing produces a high number of short lived object allocations. That might lead to some fragmentation and
+        # an unnecessarily expanded heap. Compacting ensures that the heap is as small as possible and that future
+        # allocations and garbage collections are faster
+        GC.compact unless @test_mode
+
         # Always end the progress notification even if indexing failed or else it never goes away and the user has no
         # way of dismissing it
         end_progress("indexing-progress")
@@ -912,10 +1000,10 @@ module RubyLsp
 
       return unless indexing_options
 
+      configuration = @global_state.index.configuration
+      configuration.workspace_path = @global_state.workspace_path
       # The index expects snake case configurations, but VS Code standardizes on camel case settings
-      @global_state.index.configuration.apply_config(
-        indexing_options.transform_keys { |key| key.to_s.gsub(/([A-Z])/, "_\\1").downcase },
-      )
+      configuration.apply_config(indexing_options.transform_keys { |key| key.to_s.gsub(/([A-Z])/, "_\\1").downcase })
     end
   end
 end

data/lib/ruby_lsp/setup_bundler.rb

@@ -43,13 +43,14 @@ module RubyLsp
       @gemfile_name = T.let(@gemfile&.basename&.to_s || "Gemfile", String)
 
       # Custom bundle paths
-      @custom_dir = T.let(Pathname.new(".ruby-lsp").expand_path(Dir.pwd), Pathname)
+      @custom_dir = T.let(Pathname.new(".ruby-lsp").expand_path(@project_path), Pathname)
       @custom_gemfile = T.let(@custom_dir + @gemfile_name, Pathname)
       @custom_lockfile = T.let(@custom_dir + (@lockfile&.basename || "Gemfile.lock"), Pathname)
       @lockfile_hash_path = T.let(@custom_dir + "main_lockfile_hash", Pathname)
       @last_updated_path = T.let(@custom_dir + "last_updated", Pathname)
 
       @dependencies = T.let(load_dependencies, T::Hash[String, T.untyped])
+      @rails_app = T.let(rails_app?, T::Boolean)
       @retry = T.let(false, T::Boolean)
     end
 
@@ -62,7 +63,7 @@ module RubyLsp
       # Do not set up a custom bundle if LSP dependencies are already in the Gemfile
       if @dependencies["ruby-lsp"] &&
           @dependencies["debug"] &&
-          (@dependencies["rails"] ? @dependencies["ruby-lsp-rails"] : true)
+          (@rails_app ? @dependencies["ruby-lsp-rails"] : true)
         $stderr.puts(
           "Ruby LSP> Skipping custom bundle setup since LSP dependencies are already in #{@gemfile}",
         )
@@ -148,7 +149,7 @@ module RubyLsp
         parts << 'gem "debug", require: false, group: :development, platforms: :mri'
       end
 
-      if @dependencies["rails"] && !@dependencies["ruby-lsp-rails"]
+      if @rails_app && !@dependencies["ruby-lsp-rails"]
         parts << 'gem "ruby-lsp-rails", require: false, group: :development'
       end
 
@@ -182,14 +183,14 @@ module RubyLsp
       # `.ruby-lsp` folder, which is not the user's intention. For example, if the path is configured as `vendor`, we
       # want to install it in the top level `vendor` and not `.ruby-lsp/vendor`
       path = Bundler.settings["path"]
-      expanded_path = File.expand_path(path, Dir.pwd) if path
+      expanded_path = File.expand_path(path, @project_path) if path
 
       # Use the absolute `BUNDLE_PATH` to prevent accidentally creating unwanted folders under `.ruby-lsp`
       env = {}
       env["BUNDLE_GEMFILE"] = bundle_gemfile.to_s
       env["BUNDLE_PATH"] = expanded_path if expanded_path
 
-      local_config_path = File.join(Dir.pwd, ".bundle")
+      local_config_path = File.join(@project_path, ".bundle")
       env["BUNDLE_APP_CONFIG"] = local_config_path if Dir.exist?(local_config_path)
 
       # If `ruby-lsp` and `debug` (and potentially `ruby-lsp-rails`) are already in the Gemfile, then we shouldn't try
@@ -209,7 +210,7 @@ module RubyLsp
         command << " && bundle update "
         command << "ruby-lsp " unless @dependencies["ruby-lsp"]
         command << "debug " unless @dependencies["debug"]
-        command << "ruby-lsp-rails " if @dependencies["rails"] && !@dependencies["ruby-lsp-rails"]
+        command << "ruby-lsp-rails " if @rails_app && !@dependencies["ruby-lsp-rails"]
         command << "--pre" if @experimental
         command.delete_suffix!(" ")
         command << ")"
@@ -244,7 +245,7 @@ module RubyLsp
     def should_bundle_update?
       # If `ruby-lsp`, `ruby-lsp-rails` and `debug` are in the Gemfile, then we shouldn't try to upgrade them or else it
       # will produce version control changes
-      if @dependencies["rails"]
+      if @rails_app
         return false if @dependencies.values_at("ruby-lsp", "ruby-lsp-rails", "debug").all?
 
         # If the custom lockfile doesn't include `ruby-lsp`, `ruby-lsp-rails` or `debug`, we need to run bundle install
@@ -280,5 +281,15 @@ module RubyLsp
 
       @custom_lockfile.write(content)
     end
+
+    # Detects if the project is a Rails app by looking if the superclass of the main class is `Rails::Application`
+    sig { returns(T::Boolean) }
+    def rails_app?
+      config = Pathname.new("config/application.rb").expand_path
+      application_contents = config.read(external_encoding: Encoding::UTF_8) if config.exist?
+      return false unless application_contents
+
+      /class .* < (::)?Rails::Application/.match?(application_contents)
+    end
   end
 end

data/lib/ruby_lsp/store.rb

@@ -18,7 +18,7 @@ module RubyLsp
 
     sig { void }
     def initialize
-      @state = T.let({}, T::Hash[String, Document])
+      @state = T.let({}, T::Hash[String, Document[T.untyped]])
       @supports_progress = T.let(true, T::Boolean)
       @features_configuration = T.let(
         {
@@ -33,7 +33,7 @@ module RubyLsp
       @client_name = T.let("Unknown", String)
     end
 
-    sig { params(uri: URI::Generic).returns(Document) }
+    sig { params(uri: URI::Generic).returns(Document[T.untyped]) }
     def get(uri)
       document = @state[uri.to_s]
       return document unless document.nil?
@@ -44,8 +44,11 @@ module RubyLsp
       raise NonExistingDocumentError, uri.to_s unless path
 
       ext = File.extname(path)
-      language_id = if ext == ".erb" || ext == ".rhtml"
+      language_id = case ext
+      when ".erb", ".rhtml"
         Document::LanguageId::ERB
+      when ".rbs"
+        Document::LanguageId::RBS
       else
         Document::LanguageId::Ruby
       end
@@ -66,13 +69,14 @@ module RubyLsp
       ).void
     end
     def set(uri:, source:, version:, language_id:, encoding: Encoding::UTF_8)
-      document = case language_id
+      @state[uri.to_s] = case language_id
       when Document::LanguageId::ERB
         ERBDocument.new(source: source, version: version, uri: uri, encoding: encoding)
+      when Document::LanguageId::RBS
+        RBSDocument.new(source: source, version: version, uri: uri, encoding: encoding)
       else
         RubyDocument.new(source: source, version: version, uri: uri, encoding: encoding)
       end
-      @state[uri.to_s] = document
     end
 
     sig { params(uri: URI::Generic, edits: T::Array[T::Hash[Symbol, T.untyped]], version: Integer).void }
@@ -100,7 +104,7 @@ module RubyLsp
         .params(
           uri: URI::Generic,
           request_name: String,
-          block: T.proc.params(document: Document).returns(T.type_parameter(:T)),
+          block: T.proc.params(document: Document[T.untyped]).returns(T.type_parameter(:T)),
         ).returns(T.type_parameter(:T))
     end
     def cache_fetch(uri, request_name, &block)