Add support for experimental_ksp2_psi_resolution (#1602)

* Add support for experimental_ksp2_psi_resolution

* Formatting
diff --git a/kotlin/internal/jvm/compile.bzl b/kotlin/internal/jvm/compile.bzl
index 9f5bb65..db5d96b 100644
--- a/kotlin/internal/jvm/compile.bzl
+++ b/kotlin/internal/jvm/compile.bzl
@@ -508,6 +508,10 @@
     for key, value in ksp_options.items():
         args.add("--ksp_options", "%s=%s" % (key, value))
 
+    # Toolchain-level KSP2 configuration
+    if toolchains.kt.experimental_ksp2_psi_resolution:
+        args.add("--experimental_psi_resolution", "true")
+
     # Run KSP2 via dedicated worker (separate from kotlinc worker)
     # Single action: staging + KSP2 + packaging all happen in the worker
     ctx.actions.run(
diff --git a/kotlin/internal/toolchains.bzl b/kotlin/internal/toolchains.bzl
index aff0628..9d9d0f1 100644
--- a/kotlin/internal/toolchains.bzl
+++ b/kotlin/internal/toolchains.bzl
@@ -100,6 +100,7 @@
         jacocorunner = ctx.attr.jacocorunner,
         experimental_prune_transitive_deps = ctx.attr._experimental_prune_transitive_deps[BuildSettingInfo].value,
         experimental_strict_associate_dependencies = ctx.attr._experimental_strict_associate_dependencies[BuildSettingInfo].value,
+        experimental_ksp2_psi_resolution = ctx.attr._experimental_ksp2_psi_resolution[BuildSettingInfo].value,
     )
 
     return [
@@ -312,6 +313,11 @@
             cfg = "target",
             default = Label("//third_party:empty.jdeps"),
         ),
+        "_experimental_ksp2_psi_resolution": attr.label(
+            doc = """If enabled, KSP2 uses its experimental PSI-based symbol resolution strategy
+            (KSPJvmConfig.experimentalPsiResolution) instead of the default Analysis API strategy.""",
+            default = Label("//kotlin/settings:experimental_ksp2_psi_resolution"),
+        ),
         "_experimental_prune_transitive_deps": attr.label(
             doc = """If enabled, compilation is performed against only direct dependencies.
             Transitive deps required for compilation must be explicitly added. Using
diff --git a/kotlin/settings/BUILD.bazel b/kotlin/settings/BUILD.bazel
index d311b40..a77bbd4 100644
--- a/kotlin/settings/BUILD.bazel
+++ b/kotlin/settings/BUILD.bazel
@@ -48,3 +48,11 @@
     build_setting_default = True,
     visibility = ["//visibility:public"],
 )
+
+# --@rules_kotlin//kotlin/settings:experimental_ksp2_psi_resolution=True
+# Enables KSP2's experimental PSI-based symbol resolution strategy.
+bool_flag(
+    name = "experimental_ksp2_psi_resolution",
+    build_setting_default = False,
+    visibility = ["//visibility:public"],
+)
diff --git a/kotlin/settings/BUILD.release.bazel b/kotlin/settings/BUILD.release.bazel
index 69dcf04..8d4c18c 100644
--- a/kotlin/settings/BUILD.release.bazel
+++ b/kotlin/settings/BUILD.release.bazel
@@ -40,3 +40,11 @@
     build_setting_default = True,
     visibility = ["//visibility:public"],
 )
+
+# --@rules_kotlin//kotlin/settings:experimental_ksp2_psi_resolution=True
+# Enables KSP2's experimental PSI-based symbol resolution strategy.
+bool_flag(
+    name = "experimental_ksp2_psi_resolution",
+    build_setting_default = False,
+    visibility = ["//visibility:public"],
+)
diff --git a/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/Ksp2Task.kt b/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/Ksp2Task.kt
index 531d70f..f558386 100644
--- a/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/Ksp2Task.kt
+++ b/src/main/kotlin/io/bazel/kotlin/builder/tasks/jvm/Ksp2Task.kt
@@ -66,6 +66,7 @@
       JVM_TARGET("--jvm_target"),
       JDK_HOME("--jdk_home"),
       KSP_OPTIONS("--ksp_options"),
+      EXPERIMENTAL_PSI_RESOLUTION("--experimental_psi_resolution"),
     }
 
     fun parseKspOptions(entries: List<String>): Map<String, String> =
@@ -180,6 +181,8 @@
       val kspClassLoader = URLClassLoader(processorUrls, ClassLoader.getSystemClassLoader())
 
       val processorOptions = parseKspOptions(argMap.optional(Ksp2Flags.KSP_OPTIONS) ?: emptyList())
+      val experimentalPsiResolution =
+        argMap.optionalSingle(Ksp2Flags.EXPERIMENTAL_PSI_RESOLUTION)?.toBoolean() ?: false
 
       // Load Ksp2Invoker via reflection (it's compiled against KSP2 classes)
       val invokerClass = kspClassLoader.loadClass("io.bazel.kotlin.ksp2.Ksp2Invoker")
@@ -206,6 +209,7 @@
           String::class.java, // apiVersion
           File::class.java, // jdkHome
           Map::class.java, // processorOptions
+          Boolean::class.javaPrimitiveType, // experimentalPsiResolution
           Int::class.java, // logLevel
         )
 
@@ -229,6 +233,7 @@
           argMap.optionalSingle(Ksp2Flags.API_VERSION),
           argMap.optionalSingle(Ksp2Flags.JDK_HOME)?.let { File(it) },
           processorOptions,
+          experimentalPsiResolution,
           1, // logLevel
         ) as Int
 
diff --git a/src/main/kotlin/io/bazel/kotlin/ksp2/Ksp2Invoker.kt b/src/main/kotlin/io/bazel/kotlin/ksp2/Ksp2Invoker.kt
index ba8d2ad..7ee110b 100644
--- a/src/main/kotlin/io/bazel/kotlin/ksp2/Ksp2Invoker.kt
+++ b/src/main/kotlin/io/bazel/kotlin/ksp2/Ksp2Invoker.kt
@@ -56,6 +56,7 @@
     apiVersion: String?,
     jdkHome: File?,
     processorOptions: Map<String, String> = emptyMap(),
+    experimentalPsiResolution: Boolean = false,
     logLevel: Int = 1,
   ): Int {
     // Load processors via ServiceLoader from the provided classloader
@@ -84,6 +85,7 @@
           jdkHome?.let { this.jdkHome = it }
           this.processorOptions = processorOptions
           this.mapAnnotationArgumentsInJava = true
+          this.experimentalPsiResolution = experimentalPsiResolution
         }.build()
 
     // Create logger and execute