ruby-lsp 0.20.0 → 0.21.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d43d431cf15d817d80ebb8564f90a748b871ad3c2c278482d4b1b5c831b7bcd2
-  data.tar.gz: 28ab582396723d582f7909ad277300430b6e14e6a281408be21a71f4d768930f
+  metadata.gz: a22b940469b6910a05ada8b39dbbc2ad918e26e41de0f2f287e79df46b8e46fe
+  data.tar.gz: a57e250627b82a42ba286fd9b14338c48fb04b5a2e76c889d28d93ea55f96f90
 SHA512:
-  metadata.gz: 73a87212da50592a6b201a5f21ea559de74fb12ed795428e06392a4375aa73015a103aedd56734b9b05a2f0a27fcb7366b16107ea9f1a080b3459e2eff1dcd88
-  data.tar.gz: 7ad77f3dccbabba9cb306c73ccefca4e84e9ff480bc0594d20a1c6e03727948119f7accefc522e5ed25d6de1d21fff03b18c704c7c8908b33eb92cfc5b2c665d
+  metadata.gz: 0661e457f068016e397436e76c9d48d8d4e2f718fb47268c6fe57797b1374e4669eb4c01b7a45ca26d8fe39dbc79f308797864e73aa31ba21e013ac3c2c97265
+  data.tar.gz: 709e136e94d223cb12fcb8a082d6193023fecc324aa1908f7a9dc2d9c3b612d6935cfe0362e7ba22fdd02ec3e6c1113a2a44e2a0f15230ff4698075a145d5e0d

data/README.md

@@ -29,8 +29,8 @@ Bug reports and pull requests are welcome on GitHub at https://github.com/Shopif
 be a safe, welcoming space for collaboration, and contributors are expected to adhere to the [Contributor
 Covenant](CODE_OF_CONDUCT.md) code of conduct.
 
-If you wish to contribute, see [CONTRIBUTING](CONTRIBUTING.md) for development instructions and check out our pinned
-[roadmap issue](https://github.com/Shopify/ruby-lsp/issues) for a list of tasks to get started.
+If you wish to contribute, see [Contributing](https://shopify.github.io/ruby-lsp/contributing.html) for development instructions and check out our
+[Design and roadmap](https://shopify.github.io/ruby-lsp/design-and-roadmap.html) for a list of tasks to get started.
 
 ## License
 

data/VERSION

@@ -1 +1 @@
-0.20.0
+0.21.0

data/exe/ruby-lsp

@@ -33,6 +33,10 @@ parser = OptionParser.new do |opts|
     options[:doctor] = true
   end
 
+  opts.on("--use-launcher", "[EXPERIMENTAL] Use launcher mechanism to handle missing dependencies gracefully") do
+    options[:launcher] = true
+  end
+
   opts.on("-h", "--help", "Print this help") do
     puts opts.help
     puts
@@ -54,6 +58,17 @@ end
 # using `BUNDLE_GEMFILE=.ruby-lsp/Gemfile bundle exec ruby-lsp` so that we have access to the gems that are a part of
 # the application's bundle
 if ENV["BUNDLE_GEMFILE"].nil?
+  # Substitute the current process by the launcher. RubyGems activates all dependencies of a gem's executable eagerly,
+  # but we can't have that happen because we want to invoke Bundler.setup ourselves with the composed bundle and avoid
+  # duplicate spec activation errors. Replacing the process with the launcher executable will clear the activated specs,
+  # which gives us the opportunity to control which specs are activated and enter degraded mode if any gems failed to
+  # install rather than failing to boot the server completely
+  if options[:launcher]
+    command = +File.expand_path("ruby-lsp-launcher", __dir__)
+    command << " --debug" if options[:debug]
+    exit exec(command)
+  end
+
   require_relative "../lib/ruby_lsp/setup_bundler"
 
   begin
@@ -63,12 +78,18 @@ if ENV["BUNDLE_GEMFILE"].nil?
     exit(78)
   end
 
-  exit exec(env, "bundle exec ruby-lsp #{original_args.join(" ")}")
-end
+  base_bundle = if env["BUNDLER_VERSION"]
+    "bundle _#{env["BUNDLER_VERSION"]}_"
+  else
+    "bundle"
+  end
 
-require "ruby_lsp/load_sorbet"
+  exit exec(env, "#{base_bundle} exec ruby-lsp #{original_args.join(" ")}".strip)
+end
 
 $LOAD_PATH.unshift(File.expand_path("../lib", __dir__))
+
+require "ruby_lsp/load_sorbet"
 require "ruby_lsp/internal"
 
 T::Utils.run_all_sig_blocks

data/exe/ruby-lsp-launcher

@@ -0,0 +1,127 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+# !!!!!!!
+# No gems can be required in this file until we invoke bundler setup except inside the forked process that sets up the
+# composed bundle
+# !!!!!!!
+
+setup_error = nil
+
+# Read the initialize request before even starting the server. We need to do this to figure out the workspace URI.
+# Editors are not required to spawn the language server process on the same directory as the workspace URI, so we need
+# to ensure that we're setting up the bundle in the right place
+$stdin.binmode
+headers = $stdin.gets("\r\n\r\n")
+content_length = headers[/Content-Length: (\d+)/i, 1].to_i
+raw_initialize = $stdin.read(content_length)
+
+# Compose the Ruby LSP bundle in a forked process so that we can require gems without polluting the main process
+# `$LOAD_PATH` and `Gem.loaded_specs`. Windows doesn't support forking, so we need a separate path to support it
+pid = if Gem.win_platform?
+  # Since we can't fork on Windows and spawn won't carry over the existing load paths, we need to explicitly pass that
+  # down to the child process or else requiring gems during composing the bundle will fail
+  load_path = $LOAD_PATH.flat_map do |path|
+    ["-I", File.expand_path(path)]
+  end
+
+  Process.spawn(
+    Gem.ruby,
+    *load_path,
+    File.expand_path("../lib/ruby_lsp/scripts/compose_bundle_windows.rb", __dir__),
+    raw_initialize,
+  )
+else
+  fork do
+    require_relative "../lib/ruby_lsp/scripts/compose_bundle"
+    compose(raw_initialize)
+  end
+end
+
+begin
+  # Wait until the composed Bundle is finished
+  Process.wait(pid)
+rescue Errno::ECHILD
+  # In theory, the child process can finish before we even get to the wait call, but that is not an error
+end
+
+begin
+  bundle_env_path = File.join(".ruby-lsp", "bundle_env")
+  # We can't require `bundler/setup` because that file prematurely exits the process if setup fails. However, we can't
+  # simply require bundler either because the version required might conflict with the one locked in the composed
+  # bundle. We need the composed bundle sub-process to inform us of the locked Bundler version, so that we can then
+  # activate the right spec and require the exact Bundler version required by the app
+  if File.exist?(bundle_env_path)
+    env = File.readlines(bundle_env_path).to_h { |line| line.chomp.split("=", 2) }
+    ENV.merge!(env)
+
+    if env["BUNDLER_VERSION"]
+      Gem::Specification.find_by_name("bundler", env["BUNDLER_VERSION"]).activate
+    end
+
+    require "bundler"
+    Bundler.ui.level = :silent
+    Bundler.setup
+    $stderr.puts("Composed Bundle set up successfully")
+  end
+rescue StandardError => e
+  # If installing gems failed for any reason, we don't want to exit the process prematurely. We can still provide most
+  # features in a degraded mode. We simply save the error so that we can report to the user that certain gems might be
+  # missing, but we respect the LSP life cycle
+  setup_error = e
+  $stderr.puts("Failed to set up composed Bundle\n#{e.full_message}")
+
+  # If Bundler.setup fails, we need to restore the original $LOAD_PATH so that we can still require the Ruby LSP server
+  # in degraded mode
+  $LOAD_PATH.unshift(File.expand_path("../lib", __dir__))
+ensure
+  require "fileutils"
+  FileUtils.rm(bundle_env_path) if File.exist?(bundle_env_path)
+end
+
+error_path = File.join(".ruby-lsp", "install_error")
+
+install_error = if File.exist?(error_path)
+  Marshal.load(File.read(error_path))
+end
+
+# Now that the bundle is set up, we can begin actually launching the server. Note that `Bundler.setup` will have already
+# configured the load path using the version of the Ruby LSP present in the composed bundle. Do not push any Ruby LSP
+# paths into the load path manually or we may end up requiring the wrong version of the gem
+require "ruby_lsp/load_sorbet"
+require "ruby_lsp/internal"
+
+T::Utils.run_all_sig_blocks
+
+if ARGV.include?("--debug")
+  if ["x64-mingw-ucrt", "x64-mingw32"].include?(RUBY_PLATFORM)
+    $stderr.puts "Debugging is not supported on Windows"
+  else
+    begin
+      ENV.delete("RUBY_DEBUG_IRB_CONSOLE")
+      require "debug/open_nonstop"
+    rescue LoadError
+      $stderr.puts("You need to install the debug gem to use the --debug flag")
+    end
+  end
+end
+
+# Ensure all output goes out stderr by default to allow puts/p/pp to work without specifying output device.
+$> = $stderr
+
+initialize_request = JSON.parse(raw_initialize, symbolize_names: true) if raw_initialize
+
+begin
+  RubyLsp::Server.new(
+    install_error: install_error,
+    setup_error: setup_error,
+    initialize_request: initialize_request,
+  ).start
+rescue ArgumentError
+  # If the launcher is booting an outdated version of the server, then the initializer doesn't accept a keyword splat
+  # and we already read the initialize request from the stdin pipe. In this case, we need to process the initialize
+  # request manually and then start the main loop
+  server = RubyLsp::Server.new
+  server.process_message(initialize_request)
+  server.start
+end

data/lib/ruby_indexer/lib/ruby_indexer/configuration.rb

@@ -28,7 +28,17 @@ module RubyIndexer
       @encoding = T.let(Encoding::UTF_8, Encoding)
       @excluded_gems = T.let(initial_excluded_gems, T::Array[String])
       @included_gems = T.let([], T::Array[String])
-      @excluded_patterns = T.let([File.join("**", "*_test.rb"), File.join("tmp", "**", "*")], T::Array[String])
+
+      @excluded_patterns = T.let(
+        [
+          File.join("**", "*_test.rb"),
+          File.join("node_modules", "**", "*"),
+          File.join("spec", "**", "*"),
+          File.join("test", "**", "*"),
+          File.join("tmp", "**", "*"),
+        ],
+        T::Array[String],
+      )
 
       path = Bundler.settings["path"]
       if path
@@ -56,6 +66,21 @@ module RubyIndexer
       )
     end
 
+    sig { returns(String) }
+    def merged_excluded_file_pattern
+      # This regex looks for @excluded_patterns that follow the format of "something/**/*", where
+      # "something" is one or more non-"/"
+      #
+      # Returns "/path/to/workspace/{tmp,node_modules}/**/*"
+      @excluded_patterns
+        .filter_map do |pattern|
+          next if File.absolute_path?(pattern)
+
+          pattern.match(%r{\A([^/]+)/\*\*/\*\z})&.captures&.first
+        end
+        .then { |dirs| File.join(@workspace_path, "{#{dirs.join(",")}}/**/*") }
+    end
+
     sig { returns(T::Array[IndexablePath]) }
     def indexables
       excluded_gems = @excluded_gems - @included_gems
@@ -64,21 +89,47 @@ module RubyIndexer
       # NOTE: indexing the patterns (both included and excluded) needs to happen before indexing gems, otherwise we risk
       # having duplicates if BUNDLE_PATH is set to a folder inside the project structure
 
+      flags = File::FNM_PATHNAME | File::FNM_EXTGLOB
+
+      # In order to speed up indexing, only traverse into top-level directories that are not entirely excluded.
+      # For example, if "tmp/**/*" is excluded, we don't need to traverse into "tmp" at all. However, if
+      # "vendor/bundle/**/*" is excluded, we will traverse all of "vendor" and `reject!` out all "vendor/bundle" entries
+      # later.
+      excluded_pattern = merged_excluded_file_pattern
+      included_paths = Dir.glob(File.join(@workspace_path, "*/"), flags)
+        .filter_map do |included_path|
+          next if File.fnmatch?(excluded_pattern, included_path, flags)
+
+          relative_path = included_path
+            .delete_prefix(@workspace_path)
+            .tap { |path| path.delete_prefix!("/") }
+
+          [included_path, relative_path]
+        end
+
+      indexables = T.let([], T::Array[IndexablePath])
+
       # Add user specified patterns
-      indexables = @included_patterns.flat_map do |pattern|
+      @included_patterns.each do |pattern|
         load_path_entry = T.let(nil, T.nilable(String))
 
-        Dir.glob(File.join(@workspace_path, pattern), File::FNM_PATHNAME | File::FNM_EXTGLOB).map! do |path|
-          path = File.expand_path(path)
-          # All entries for the same pattern match the same $LOAD_PATH entry. Since searching the $LOAD_PATH for every
-          # entry is expensive, we memoize it until we find a path that doesn't belong to that $LOAD_PATH. This happens
-          # on repositories that define multiple gems, like Rails. All frameworks are defined inside the current
-          # workspace directory, but each one of them belongs to a different $LOAD_PATH entry
-          if load_path_entry.nil? || !path.start_with?(load_path_entry)
-            load_path_entry = $LOAD_PATH.find { |load_path| path.start_with?(load_path) }
-          end
+        included_paths.each do |included_path, relative_path|
+          relative_pattern = pattern.delete_prefix(File.join(relative_path, "/"))
+
+          next unless pattern.start_with?("**") || pattern.start_with?(relative_path)
 
-          IndexablePath.new(load_path_entry, path)
+          Dir.glob(File.join(included_path, relative_pattern), flags).each do |path|
+            path = File.expand_path(path)
+            # All entries for the same pattern match the same $LOAD_PATH entry. Since searching the $LOAD_PATH for every
+            # entry is expensive, we memoize it until we find a path that doesn't belong to that $LOAD_PATH. This
+            # happens on repositories that define multiple gems, like Rails. All frameworks are defined inside the
+            # current workspace directory, but each one of them belongs to a different $LOAD_PATH entry
+            if load_path_entry.nil? || !path.start_with?(load_path_entry)
+              load_path_entry = $LOAD_PATH.find { |load_path| path.start_with?(load_path) }
+            end
+
+            indexables << IndexablePath.new(load_path_entry, path)
+          end
         end
       end
 

data/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -312,12 +312,16 @@ module RubyIndexer
         @visibility_stack.push(Entry::Visibility::PROTECTED)
       when :private
         @visibility_stack.push(Entry::Visibility::PRIVATE)
+      when :module_function
+        handle_module_function(node)
       end
 
       @enhancements.each do |enhancement|
-        enhancement.on_call_node(@index, @owner_stack.last, node, @file_path, @code_units_cache)
+        enhancement.on_call_node_enter(@owner_stack.last, node, @file_path, @code_units_cache)
       rescue StandardError => e
-        @indexing_errors << "Indexing error in #{@file_path} with '#{enhancement.class.name}' enhancement: #{e.message}"
+        @indexing_errors << <<~MSG
+          Indexing error in #{@file_path} with '#{enhancement.class.name}' on call node enter enhancement: #{e.message}
+        MSG
       end
     end
 
@@ -332,6 +336,14 @@ module RubyIndexer
           @visibility_stack.pop
         end
       end
+
+      @enhancements.each do |enhancement|
+        enhancement.on_call_node_leave(@owner_stack.last, node, @file_path, @code_units_cache)
+      rescue StandardError => e
+        @indexing_errors << <<~MSG
+          Indexing error in #{@file_path} with '#{enhancement.class.name}' on call node leave enhancement: #{e.message}
+        MSG
+      end
     end
 
     sig { params(node: Prism::DefNode).void }
@@ -751,6 +763,48 @@ module RubyIndexer
       end
     end
 
+    sig { params(node: Prism::CallNode).void }
+    def handle_module_function(node)
+      arguments_node = node.arguments
+      return unless arguments_node
+
+      owner_name = @owner_stack.last&.name
+      return unless owner_name
+
+      arguments_node.arguments.each do |argument|
+        method_name = case argument
+        when Prism::StringNode
+          argument.content
+        when Prism::SymbolNode
+          argument.value
+        end
+        next unless method_name
+
+        entries = @index.resolve_method(method_name, owner_name)
+        next unless entries
+
+        entries.each do |entry|
+          entry_owner_name = entry.owner&.name
+          next unless entry_owner_name
+
+          entry.visibility = Entry::Visibility::PRIVATE
+
+          singleton = @index.existing_or_new_singleton_class(entry_owner_name)
+          location = Location.from_prism_location(argument.location, @code_units_cache)
+          @index.add(Entry::Method.new(
+            method_name,
+            @file_path,
+            location,
+            location,
+            collect_comments(node)&.concat(entry.comments),
+            entry.signatures,
+            Entry::Visibility::PUBLIC,
+            singleton,
+          ))
+        end
+      end
+    end
+
     sig { returns(Entry::Visibility) }
     def current_visibility
       T.must(@visibility_stack.last)

data/lib/ruby_indexer/lib/ruby_indexer/enhancement.rb

@@ -2,20 +2,35 @@
 # frozen_string_literal: true
 
 module RubyIndexer
-  module Enhancement
+  class Enhancement
     extend T::Sig
     extend T::Helpers
 
-    interface!
+    abstract!
 
-    requires_ancestor { Object }
+    sig { params(index: Index).void }
+    def initialize(index)
+      @index = index
+    end
 
     # The `on_extend` indexing enhancement is invoked whenever an extend is encountered in the code. It can be used to
     # register for an included callback, similar to what `ActiveSupport::Concern` does in order to auto-extend the
     # `ClassMethods` modules
     sig do
-      abstract.params(
-        index: Index,
+      overridable.params(
+        owner: T.nilable(Entry::Namespace),
+        node: Prism::CallNode,
+        file_path: String,
+        code_units_cache: T.any(
+          T.proc.params(arg0: Integer).returns(Integer),
+          Prism::CodeUnitsCache,
+        ),
+      ).void
+    end
+    def on_call_node_enter(owner, node, file_path, code_units_cache); end
+
+    sig do
+      overridable.params(
         owner: T.nilable(Entry::Namespace),
         node: Prism::CallNode,
         file_path: String,
@@ -25,6 +40,6 @@ module RubyIndexer
         ),
       ).void
     end
-    def on_call_node(index, owner, node, file_path, code_units_cache); end
+    def on_call_node_leave(owner, node, file_path, code_units_cache); end
   end
 end

data/lib/ruby_indexer/lib/ruby_indexer/entry.rb

@@ -646,7 +646,7 @@ module RubyIndexer
         (positionals.empty? && forwarding_arguments.any?) ||
           (
             # Check if positional arguments match. This includes required, optional, rest arguments. We also need to
-            # verify if there's a trailing forwading argument, like `def foo(a, ...); end`
+            # verify if there's a trailing forwarding argument, like `def foo(a, ...); end`
             positional_arguments_match?(positionals, forwarding_arguments, keyword_args, min_pos, max_pos) &&
             # If the positional arguments match, we move on to checking keyword, optional keyword and keyword rest
             # arguments. If there's a forward argument, then it will always match. If the method accepts a keyword rest

data/lib/ruby_indexer/lib/ruby_indexer/index.rb

@@ -784,7 +784,7 @@ module RubyIndexer
       singleton_levels
     )
       # Find the first class entry that has a parent class. Notice that if the developer makes a mistake and inherits
-      # from two diffent classes in different files, we simply ignore it
+      # from two different classes in different files, we simply ignore it
       superclass = T.cast(
         if singleton_levels > 0
           self[attached_class_name]&.find { |n| n.is_a?(Entry::Class) && n.parent_class }
@@ -974,35 +974,29 @@ module RubyIndexer
       []
     end
 
-    # Removes redudancy from a constant reference's full name. For example, if we find a reference to `A::B::Foo` inside
-    # of the ["A", "B"] nesting, then we should not concatenate the nesting with the name or else we'll end up with
-    # `A::B::A::B::Foo`. This method will remove any redundant parts from the final name based on the reference and the
-    # nesting
+    # Removes redundancy from a constant reference's full name. For example, if we find a reference to `A::B::Foo`
+    # inside of the ["A", "B"] nesting, then we should not concatenate the nesting with the name or else we'll end up
+    # with `A::B::A::B::Foo`. This method will remove any redundant parts from the final name based on the reference and
+    # the nesting
     sig { params(name: String, nesting: T::Array[String]).returns(String) }
     def build_non_redundant_full_name(name, nesting)
+      # If there's no nesting, then we can just return the name as is
       return name if nesting.empty?
 
-      namespace = nesting.join("::")
-
       # If the name is not qualified, we can just concatenate the nesting and the name
-      return "#{namespace}::#{name}" unless name.include?("::")
+      return "#{nesting.join("::")}::#{name}" unless name.include?("::")
 
       name_parts = name.split("::")
+      first_redundant_part = nesting.index(name_parts[0])
 
-      # Find the first part of the name that is not in the nesting
-      index = name_parts.index { |part| !nesting.include?(part) }
+      # If there are no redundant parts between the name and the nesting, then the full name is both combined
+      return "#{nesting.join("::")}::#{name}" unless first_redundant_part
 
-      if index.nil?
-        # All parts of the nesting are redundant because they are already present in the name. We can return the name
-        # directly
-        name
-      elsif index == 0
-        # No parts of the nesting are in the name, we can concatenate the namespace and the name
-        "#{namespace}::#{name}"
-      else
-        # The name includes some parts of the nesting. We need to remove the redundant parts
-        "#{namespace}::#{T.must(name_parts[index..-1]).join("::")}"
-      end
+      # Otherwise, push all of the leading parts of the nesting that aren't redundant into the name. For example, if we
+      # have a reference to `Foo::Bar` inside the `[Namespace, Foo]` nesting, then only the `Foo` part is redundant, but
+      # we still need to include the `Namespace` part
+      T.unsafe(name_parts).unshift(*nesting[0...first_redundant_part])
+      name_parts.join("::")
     end
 
     sig do

data/lib/ruby_indexer/test/classes_and_modules_test.rb

@@ -72,7 +72,7 @@ module RubyIndexer
       assert_entry("self::Bar", Entry::Class, "/fake/path/foo.rb:0-0:1-3")
     end
 
-    def test_dynamically_namespaced_class_doesnt_affect_other_classes
+    def test_dynamically_namespaced_class_does_not_affect_other_classes
       index(<<~RUBY)
         class Foo
           class self::Bar
@@ -143,7 +143,7 @@ module RubyIndexer
       assert_entry("self::Bar", Entry::Module, "/fake/path/foo.rb:0-0:1-3")
     end
 
-    def test_dynamically_namespaced_module_doesnt_affect_other_modules
+    def test_dynamically_namespaced_module_does_not_affect_other_modules
       index(<<~RUBY)
         module Foo
           class self::Bar

data/lib/ruby_indexer/test/enhancements_test.rb

@@ -6,10 +6,8 @@ require_relative "test_case"
 module RubyIndexer
   class EnhancementTest < TestCase
     def test_enhancing_indexing_included_hook
-      enhancement_class = Class.new do
-        include Enhancement
-
-        def on_call_node(index, owner, node, file_path, code_units_cache)
+      enhancement_class = Class.new(Enhancement) do
+        def on_call_node_enter(owner, node, file_path, code_units_cache)
           return unless owner
           return unless node.name == :extend
 
@@ -24,7 +22,7 @@ module RubyIndexer
             module_name = node.full_name
             next unless module_name == "ActiveSupport::Concern"
 
-            index.register_included_hook(owner.name) do |index, base|
+            @index.register_included_hook(owner.name) do |index, base|
               class_methods_name = "#{owner.name}::ClassMethods"
 
               if index.indexed?(class_methods_name)
@@ -33,7 +31,7 @@ module RubyIndexer
               end
             end
 
-            index.add(Entry::Method.new(
+            @index.add(Entry::Method.new(
               "new_method",
               file_path,
               location,
@@ -50,7 +48,7 @@ module RubyIndexer
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
       index(<<~RUBY)
         module ActiveSupport
           module Concern
@@ -98,10 +96,8 @@ module RubyIndexer
     end
 
     def test_enhancing_indexing_configuration_dsl
-      enhancement_class = Class.new do
-        include Enhancement
-
-        def on_call_node(index, owner, node, file_path, code_units_cache)
+      enhancement_class = Class.new(Enhancement) do
+        def on_call_node_enter(owner, node, file_path, code_units_cache)
           return unless owner
 
           name = node.name
@@ -115,7 +111,7 @@ module RubyIndexer
 
           location = Location.from_prism_location(association_name.location, code_units_cache)
 
-          index.add(Entry::Method.new(
+          @index.add(Entry::Method.new(
             T.must(association_name.value),
             file_path,
             location,
@@ -128,7 +124,7 @@ module RubyIndexer
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
       index(<<~RUBY)
         module ActiveSupport
           module Concern
@@ -160,11 +156,44 @@ module RubyIndexer
       assert_entry("posts", Entry::Method, "/fake/path/foo.rb:23-11:23-17")
     end
 
-    def test_error_handling_in_enhancement
-      enhancement_class = Class.new do
-        include Enhancement
+    def test_error_handling_in_on_call_node_enter_enhancement
+      enhancement_class = Class.new(Enhancement) do
+        def on_call_node_enter(owner, node, file_path, code_units_cache)
+          raise "Error"
+        end
+
+        class << self
+          def name
+            "TestEnhancement"
+          end
+        end
+      end
+
+      @index.register_enhancement(enhancement_class.new(@index))
+
+      _stdout, stderr = capture_io do
+        index(<<~RUBY)
+          module ActiveSupport
+            module Concern
+              def self.extended(base)
+                base.class_eval("def new_method(a); end")
+              end
+            end
+          end
+        RUBY
+      end
+
+      assert_match(
+        %r{Indexing error in /fake/path/foo\.rb with 'TestEnhancement' on call node enter enhancement},
+        stderr,
+      )
+      # The module should still be indexed
+      assert_entry("ActiveSupport::Concern", Entry::Module, "/fake/path/foo.rb:1-2:5-5")
+    end
 
-        def on_call_node(index, owner, node, file_path, code_units_cache)
+    def test_error_handling_in_on_call_node_leave_enhancement
+      enhancement_class = Class.new(Enhancement) do
+        def on_call_node_leave(owner, node, file_path, code_units_cache)
           raise "Error"
         end
 
@@ -175,7 +204,7 @@ module RubyIndexer
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
 
       _stdout, stderr = capture_io do
         index(<<~RUBY)
@@ -189,7 +218,10 @@ module RubyIndexer
         RUBY
       end
 
-      assert_match(%r{Indexing error in /fake/path/foo\.rb with 'TestEnhancement' enhancement}, stderr)
+      assert_match(
+        %r{Indexing error in /fake/path/foo\.rb with 'TestEnhancement' on call node leave enhancement},
+        stderr,
+      )
       # The module should still be indexed
       assert_entry("ActiveSupport::Concern", Entry::Module, "/fake/path/foo.rb:1-2:5-5")
     end