handle exception instead of crashing (#18895)

Fixes: #18641

PiperOrigin-RevId: 546092727
Change-Id: I66365813a40390ec61a1a0565b06655b3ca50fcd

Author: SalmaSamy

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java

@@ -771,9 +771,15 @@ public String getName() {
     public Object call(StarlarkThread thread, Tuple args, Dict<String, Object> kwargs)
         throws EvalException, InterruptedException {
       if (!args.isEmpty()) {
-        throw new EvalException("unexpected positional arguments");
+        throw new EvalException("Unexpected positional arguments");
+      }
+      try {
+        BazelStarlarkContext.from(thread).checkLoadingPhase(getName());
+      } catch (IllegalStateException e) {
+        throw new EvalException(
+            "A rule can only be instantiated in a BUILD file, or a macro "
+                + "invoked from a BUILD file");
       }
-      BazelStarlarkContext.from(thread).checkLoadingPhase(getName());
       if (ruleClass == null) {
         throw new EvalException("Invalid rule class hasn't been exported by a bzl file");
       }

src/test/py/bazel/bzlmod/bazel_module_test.py

@@ -635,6 +635,37 @@ def testWorkspaceToolchainRegistrationWithPlatformsConstraint(self):
     with open(self.Path('bazel-bin/my_consumer'), 'r') as f:
       self.assertEqual(f.read().strip(), 'my_value = Hello, Bzlmod!')
 
+  def testModuleExtensionWithRuleError(self):
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'ext = use_extension("extensions.bzl", "ext")',
+            'use_repo(ext, "ext")',
+        ],
+    )
+    self.ScratchFile('BUILD')
+    self.ScratchFile(
+        'extensions.bzl',
+        [
+            'def _rule_impl(ctx):',
+            '  print("RULE CALLED")',
+            'init_rule = rule(_rule_impl)',
+            'def ext_impl(module_ctx):',
+            '  init_rule()',
+            'ext = module_extension(implementation = ext_impl,)',
+        ],
+    )
+    exit_code, _, stderr = self.RunBazel(
+        ['build', '--nobuild', '@ext//:all'],
+        allow_failure=True,
+    )
+    self.AssertExitCode(exit_code, 48, stderr)
+    self.assertIn(
+        'Error in init_rule: A rule can only be instantiated in a BUILD file, '
+        'or a macro invoked from a BUILD file',
+        stderr,
+    )
+
 
 if __name__ == '__main__':
   unittest.main()