move BSP QueryRule registration out of core rules (#21311)

The refactor in #20913 to move
BSP out of the core rules (via the new `AuxiliaryGoal` support) did not
move the registration of BSP-related `QueryRule`s out of the core rules.

This PR remedies that oversight by moving the `QueryRule` registration
to the applicable BSP backends.

src/python/pants/backend/java/bsp/rules.py

@@ -9,6 +9,7 @@
 from pants.bsp.protocol import BSPHandlerMapping
 from pants.bsp.spec.base import BuildTargetIdentifier
 from pants.bsp.util_rules.lifecycle import BSPLanguageSupport
+from pants.bsp.util_rules.queries import compute_handler_query_rules
 from pants.bsp.util_rules.targets import (
     BSPBuildTargetsMetadataRequest,
     BSPBuildTargetsMetadataResult,
@@ -149,7 +150,7 @@ async def bsp_java_resources_request(
 
 
 def rules():
-    return (
+    base_rules = (
         *collect_rules(),
         *jvm_compile_rules(),
         *jvm_resources_rules(),
@@ -159,3 +160,4 @@ def rules():
         UnionRule(BSPCompileRequest, JavaBSPCompileRequest),
         UnionRule(BSPResourcesRequest, JavaBSPResourcesRequest),
     )
+    return (*base_rules, *compute_handler_query_rules(base_rules))

src/python/pants/backend/scala/bsp/rules.py

@@ -38,6 +38,7 @@
 from pants.bsp.spec.base import BuildTargetIdentifier
 from pants.bsp.spec.targets import DependencyModule
 from pants.bsp.util_rules.lifecycle import BSPLanguageSupport
+from pants.bsp.util_rules.queries import compute_handler_query_rules
 from pants.bsp.util_rules.targets import (
     BSPBuildTargetsMetadataRequest,
     BSPBuildTargetsMetadataResult,
@@ -537,7 +538,7 @@ async def bsp_scala_resources_request(
 
 
 def rules():
-    return (
+    base_rules = (
         *collect_rules(),
         *jvm_compile_rules(),
         *jvm_resources_rules(),
@@ -550,3 +551,4 @@ def rules():
         UnionRule(BSPResourcesRequest, ScalaBSPResourcesRequest),
         UnionRule(BSPDependencyModulesRequest, ScalaBSPDependencyModulesRequest),
     )
+    return (*base_rules, *compute_handler_query_rules(base_rules))

src/python/pants/bsp/rules.py

@@ -4,6 +4,7 @@
 
 from pants.bsp.context import BSPContext
 from pants.bsp.util_rules import compile, lifecycle, resources, targets
+from pants.bsp.util_rules.queries import compute_handler_query_rules
 from pants.engine.internals.session import SessionValues
 from pants.engine.rules import collect_rules, rule
 
@@ -14,10 +15,12 @@ async def bsp_context(session_values: SessionValues) -> BSPContext:
 
 
 def rules():
-    return (
+    base_rules = (
         *collect_rules(),
         *compile.rules(),
         *lifecycle.rules(),
         *resources.rules(),
         *targets.rules(),
     )
+
+    return (*base_rules, *compute_handler_query_rules(base_rules))

src/python/pants/bsp/util_rules/queries.py

@@ -0,0 +1,31 @@
+# Copyright 2024 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+from __future__ import annotations
+
+from typing import Iterable
+
+from pants.bsp.protocol import BSPHandlerMapping
+from pants.engine.environment import EnvironmentName
+from pants.engine.fs import Workspace
+from pants.engine.rules import QueryRule, Rule
+from pants.engine.unions import UnionRule
+
+
+# Compute QueryRule's for each handler request/response pair.
+# Note: These are necessary because the BSP support is an auxiliary goal that makes
+# synchronous requests into the engine.
+def compute_handler_query_rules(
+    rules: Iterable[Rule | UnionRule | QueryRule],
+) -> tuple[QueryRule, ...]:
+    queries: list[QueryRule] = []
+
+    for rule in rules:
+        if isinstance(rule, UnionRule):
+            if rule.union_base == BSPHandlerMapping:
+                impl = rule.union_member
+                assert issubclass(impl, BSPHandlerMapping)
+                queries.append(
+                    QueryRule(impl.response_type, (impl.request_type, Workspace, EnvironmentName))
+                )
+
+    return tuple(queries)

src/python/pants/init/engine_initializer.py

@@ -12,7 +12,6 @@
 from pants.base.build_root import BuildRoot
 from pants.base.exiter import PANTS_SUCCEEDED_EXIT_CODE
 from pants.base.specs import Specs
-from pants.bsp.protocol import BSPHandlerMapping
 from pants.build_graph.build_configuration import BuildConfiguration
 from pants.core.util_rules import environments, system_binaries
 from pants.core.util_rules.environments import determine_bootstrap_environment
@@ -324,13 +323,6 @@ def current_executing_goals(session_values: SessionValues) -> CurrentExecutingGo
                     )
                     for goal_type in goal_map.values()
                 ),
-                # Install queries for each request/response pair used by the BSP support.
-                # Note: These are necessary because the BSP support is a built-in goal that makes
-                # synchronous requests into the engine.
-                *(
-                    QueryRule(impl.response_type, (impl.request_type, Workspace, EnvironmentName))
-                    for impl in union_membership.get(BSPHandlerMapping)
-                ),
                 QueryRule(Snapshot, [PathGlobs]),  # Used by the SchedulerService.
             )
         )