From a101cbdb6be46dfc4c1d62c872e26bf4629d2d2a Mon Sep 17 00:00:00 2001
From: Vinicius Stock <vinicius.stock@shopify.com>
Date: Thu, 26 Feb 2026 15:34:10 -0500
Subject: [PATCH] Include RBS paths when indexing workspace

---
 Gemfile              |  1 +
 Gemfile.lock         | 10 ++++++++--
 lib/rubydex/graph.rb | 17 +++++++++++++++++
 test/graph_test.rb   | 29 +++++++++++++++++++++++++++++
 4 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/Gemfile b/Gemfile
index 76b8c9a2..c3c05c70 100644
--- a/Gemfile
+++ b/Gemfile
@@ -11,6 +11,7 @@ gem "minitest"
 gem "rubocop"
 gem "rubocop-shopify"
 gem "extconf_compile_commands_json"
+gem "rbs"
 
 # Gems that aren't supported on Windows
 platforms :ruby do
diff --git a/Gemfile.lock b/Gemfile.lock
index 02821f87..026e7144 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -11,6 +11,7 @@ GEM
     json (2.18.0)
     language_server-protocol (3.17.0.5)
     lint_roller (1.1.0)
+    logger (1.7.0)
     mini_portile2 (2.8.9)
     minitest (6.0.1)
       prism (~> 1.5)
@@ -29,6 +30,9 @@ GEM
     rake (13.3.1)
     rake-compiler (1.3.0)
       rake
+    rbs (3.10.3)
+      logger
+      tsort
     regexp_parser (2.11.3)
     rubocop (1.81.7)
       json (~> 2.3)
@@ -49,12 +53,13 @@ GEM
     ruby-progressbar (1.13.0)
     ruby_memcheck (3.0.1)
       nokogiri
+    tsort (0.2.0)
     unicode-display_width (3.2.0)
       unicode-emoji (~> 4.1)
     unicode-emoji (4.2.0)
 
 PLATFORMS
-  arm64-darwin-23
+  arm64-darwin-24
   ruby
 
 DEPENDENCIES
@@ -62,10 +67,11 @@ DEPENDENCIES
   minitest
   rake (~> 13.3)
   rake-compiler
+  rbs
   rubocop
   rubocop-shopify
   ruby_memcheck
   rubydex!
 
 BUNDLED WITH
-  4.0.3
+  4.0.7
diff --git a/lib/rubydex/graph.rb b/lib/rubydex/graph.rb
index 672f9ae9..cf4d35c1 100644
--- a/lib/rubydex/graph.rb
+++ b/lib/rubydex/graph.rb
@@ -45,6 +45,7 @@ def workspace_paths
       end
 
       add_workspace_dependency_paths(paths)
+      add_core_rbs_definition_paths(paths)
       paths.uniq!
       paths
     end
@@ -71,5 +72,21 @@ def add_workspace_dependency_paths(paths)
         nil
       end
     end
+
+    # Searches for the latest installation of the `rbs` gem and adds the paths for the core and stdlib RBS definitions
+    # to the list of paths. This method does not require `rbs` to be a part of the bundle. It searches for whatever
+    # latest installation of `rbs` exists in the system and fails silently if we can't find one
+    #
+    #: (Array[String]) -> void
+    def add_core_rbs_definition_paths(paths)
+      rbs_gem_path = Gem.path
+        .flat_map { |path| Dir.glob(File.join(path, "gems", "rbs-[0-9]*/")) }
+        .max_by { |path| Gem::Version.new(File.basename(path).delete_prefix("rbs-")) }
+
+      return unless rbs_gem_path
+
+      paths << File.join(rbs_gem_path, "core")
+      paths << File.join(rbs_gem_path, "stdlib")
+    end
   end
 end
diff --git a/test/graph_test.rb b/test/graph_test.rb
index ccd54b3c..9a8131a1 100644
--- a/test/graph_test.rb
+++ b/test/graph_test.rb
@@ -560,6 +560,35 @@ def test_workspace_paths
     end
   end
 
+  def test_index_workspace_includes_rbs_core_definitions
+    graph = Rubydex::Graph.new
+    graph.index_workspace
+    graph.resolve
+
+    ["Kernel", "Object", "BasicObject", "Integer"].each do |core_namespace|
+      rbs_kernel = graph[core_namespace].definitions.find do |definition|
+        uri = URI(definition.location.uri)
+        File.extname(uri.path) == ".rbs"
+      end
+      assert(rbs_kernel, "Expected to find RBS definition for `#{core_namespace}` in the graph")
+    end
+  end
+
+  def test_index_workspace_includes_user_defined_rbs_files
+    with_context do |context|
+      context.write!("sig/foo.rbs", <<~RBS)
+        class Foo
+        end
+      RBS
+
+      graph = Rubydex::Graph.new(workspace_path: context.absolute_path)
+      graph.index_workspace
+      graph.resolve
+
+      assert_equal("Foo", graph["Foo"].name)
+    end
+  end
+
   private
 
   def assert_diagnostics(expected, actual)