bugfix: Allow run in Mill BSP and remove non working test lenses

Previously, we would show test code lense for Mill BSP, which did not work and not show run lenses which could work with newer Mill versions. Now, we only show run code lense, which can change id debug is supported by mill or we implement an alternative way of running tests.

metals/src/main/scala/scala/meta/internal/metals/BuildServerConnection.scala

@@ -72,6 +72,8 @@ class BuildServerConnection private (
 
   // the name is set before when establishing conenction
   def name: String = initialConnection.socketConnection.serverName
+  private def capabilities: BuildServerCapabilities =
+    initialConnection.capabilities
 
   def isBloop: Boolean = name == BloopServers.name
 
@@ -83,6 +85,13 @@ class BuildServerConnection private (
 
   def isAmmonite: Boolean = name == Ammonite.name
 
+  def isDebuggingProvider: Boolean =
+    Option(capabilities.getDebugProvider())
+      .exists(_.getLanguageIds().contains("scala"))
+
+  def isJvmEnvironmentSupported: Boolean =
+    capabilities.getJvmRunEnvironmentProvider()
+
   /* Currently only Bloop and sbt support running single test cases
    * and ScalaCLI uses Bloop underneath.
    */

metals/src/main/scala/scala/meta/internal/metals/MetalsLspService.scala

@@ -1575,13 +1575,15 @@ class MetalsLspService(
   override def codeLens(
       params: CodeLensParams
   ): CompletableFuture[util.List[CodeLens]] =
-    CancelTokens { _ =>
-      timerProvider.timedThunk(
-        "code lens generation",
-        thresholdMillis = 1.second.toMillis,
-      ) {
-        val path = params.getTextDocument.getUri.toAbsolutePath
-        codeLensProvider.findLenses(path).toList.asJava
+    CancelTokens.future { _ =>
+      buildServerPromise.future.map { _ =>
+        timerProvider.timedThunk(
+          "code lens generation",
+          thresholdMillis = 1.second.toMillis,
+        ) {
+          val path = params.getTextDocument.getUri.toAbsolutePath
+          codeLensProvider.findLenses(path).toList.asJava
+        }
       }
     }
 

metals/src/main/scala/scala/meta/internal/metals/codelenses/RunTestCodeLens.scala

@@ -73,17 +73,26 @@ final class RunTestCodeLens(
       // most of the bsp servers such as bloop and sbt might not support it.
     } yield {
       val classes = buildTargetClasses.classesOf(buildTargetId)
+      // sbt doesn't declare debugging provider
+      def buildServerCanDebug =
+        connection.isDebuggingProvider || connection.isSbt
+
       if (connection.isScalaCLI && path.isAmmoniteScript) {
-        scalaCliCodeLenses(textDocument, buildTargetId, classes, distance)
-      } else if (
-        buildTarget.getCapabilities.getCanDebug || connection.isBloop || connection.isSbt
-      ) {
+        scalaCliCodeLenses(
+          textDocument,
+          buildTargetId,
+          classes,
+          distance,
+          buildServerCanDebug,
+        )
+      } else if (buildServerCanDebug || clientConfig.isRunProvider()) {
         codeLenses(
           textDocument,
           buildTargetId,
           classes,
           distance,
           path,
+          buildServerCanDebug,
         )
       } else { Nil }
 
@@ -132,6 +141,7 @@ final class RunTestCodeLens(
       occurence: SymbolOccurrence,
       textDocument: TextDocument,
       target: BuildTargetIdentifier,
+      buildServerCanDebug: Boolean,
   ): Seq[l.Command] = {
     if (occurence.symbol.endsWith("#main().")) {
       textDocument.symbols
@@ -148,6 +158,7 @@ final class RunTestCodeLens(
                 Nil.asJava,
                 Nil.asJava,
               ),
+              buildServerCanDebug,
             )
           else
             Nil
@@ -164,6 +175,7 @@ final class RunTestCodeLens(
       classes: BuildTargetClasses.Classes,
       distance: TokenEditDistance,
       path: AbsolutePath,
+      buildServerCanDebug: Boolean,
   ): Seq[l.CodeLens] = {
     for {
       occurrence <- textDocument.occurrences
@@ -172,23 +184,25 @@ final class RunTestCodeLens(
       commands = {
         val main = classes.mainClasses
           .get(symbol)
-          .map(mainCommand(target, _))
+          .map(mainCommand(target, _, buildServerCanDebug))
           .getOrElse(Nil)
         val tests =
           // Currently tests can only be run via DAP
-          if (clientConfig.isDebuggingProvider())
+          if (clientConfig.isDebuggingProvider() && buildServerCanDebug)
             testClasses(target, classes, symbol)
           else Nil
         val fromAnnot = DebugProvider
           .mainFromAnnotation(occurrence, textDocument)
           .flatMap { symbol =>
             classes.mainClasses
               .get(symbol)
-              .map(mainCommand(target, _))
+              .map(mainCommand(target, _, buildServerCanDebug))
           }
           .getOrElse(Nil)
         val javaMains =
-          if (path.isJava) javaLenses(occurrence, textDocument, target) else Nil
+          if (path.isJava)
+            javaLenses(occurrence, textDocument, target, buildServerCanDebug)
+          else Nil
         main ++ tests ++ fromAnnot ++ javaMains
       }
       if commands.nonEmpty
@@ -209,6 +223,7 @@ final class RunTestCodeLens(
       target: BuildTargetIdentifier,
       classes: BuildTargetClasses.Classes,
       distance: TokenEditDistance,
+      buildServerCanDebug: Boolean,
   ): Seq[l.CodeLens] = {
     val scriptFileName = textDocument.uri.stripSuffix(".sc")
 
@@ -218,15 +233,15 @@ final class RunTestCodeLens(
     val main =
       classes.mainClasses
         .get(expectedMainClass)
-        .map(mainCommand(target, _))
+        .map(mainCommand(target, _, buildServerCanDebug))
         .getOrElse(Nil)
 
     val fromAnnotations = textDocument.occurrences.flatMap { occ =>
       for {
         sym <- DebugProvider.mainFromAnnotation(occ, textDocument)
         cls <- classes.mainClasses.get(sym)
         range <- occurrenceRange(occ, distance)
-      } yield mainCommand(target, cls).map { cmd =>
+      } yield mainCommand(target, cls, buildServerCanDebug).map { cmd =>
         new l.CodeLens(range, cmd, null)
       }
     }.flatten
@@ -276,26 +291,30 @@ final class RunTestCodeLens(
   private def mainCommand(
       target: b.BuildTargetIdentifier,
       main: b.ScalaMainClass,
+      buildServerCanDebug: Boolean,
   ): List[l.Command] = {
-    val data = buildTargetClasses.jvmRunEnvironment
+    val (data, shellCommandAdded) = buildTargetClasses.jvmRunEnvironment
       .get(target)
       .zip(userConfig().usedJavaBinary) match {
       case None =>
-        main.toJson
+        (main.toJson, false)
       case Some((env, javaHome)) =>
-        ExtendedScalaMainClass(main, env, javaHome, workspace).toJson
+        (ExtendedScalaMainClass(main, env, javaHome, workspace).toJson, true)
     }
     val params = {
       val dataKind = b.DebugSessionParamsDataKind.SCALA_MAIN_CLASS
       sessionParams(target, dataKind, data)
     }
 
-    if (clientConfig.isDebuggingProvider())
+    if (clientConfig.isDebuggingProvider() && buildServerCanDebug)
       List(
         command("run", StartRunSession, params),
         command("debug", StartDebugSession, params),
       )
-    else List(command("run", StartRunSession, params))
+    // run provider needs shell command to run currently, we don't support pure run inside metals
+    else if (shellCommandAdded && clientConfig.isRunProvider())
+      List(command("run", StartRunSession, params))
+    else Nil
   }
 
   private def sessionParams(

metals/src/main/scala/scala/meta/internal/metals/debug/BuildTargetClasses.scala

@@ -76,9 +76,13 @@ final class BuildTargetClasses(
             .mainClasses(new b.ScalaMainClassesParams(targetsList))
             .map(cacheMainClasses(classes, _))
 
-          val updateTestClasses = connection
-            .testClasses(new b.ScalaTestClassesParams(targetsList))
-            .map(cacheTestClasses(classes, _))
+          // Currently tests are only run using DAP
+          val updateTestClasses =
+            if (connection.isDebuggingProvider || connection.isSbt)
+              connection
+                .testClasses(new b.ScalaTestClassesParams(targetsList))
+                .map(cacheTestClasses(classes, _))
+            else Future.unit
 
           val jvmRunEnvironment = connection
             .jvmRunEnvironment(new b.JvmRunEnvironmentParams(targetsList))

project/V.scala

@@ -32,7 +32,7 @@ object V {
   val kindProjector = "0.13.2"
   val lsp4jV = "0.20.1"
   val mavenBloop = "2.0.0"
-  val mill = "0.11.1"
+  val mill = "0.11.2"
   val mdoc = "2.3.7"
   val munit = "1.0.0-M8"
   val pprint = "0.7.3"

tests/slow/src/test/scala/tests/mill/MillServerCodeLensSuite.scala

@@ -0,0 +1,69 @@
+package tests.mill
+
+import scala.concurrent.duration.Duration
+
+import scala.meta.internal.metals.ServerCommands
+import scala.meta.internal.metals.{BuildInfo => V}
+
+import tests.BaseCodeLensLspSuite
+import tests.MillBuildLayout
+import tests.MillServerInitializer
+
+class MillServerCodeLensSuite
+    extends BaseCodeLensLspSuite("mill-server-lenses", MillServerInitializer) {
+
+  override def munitTimeout: Duration = Duration("4min")
+
+  test("run-mill-lens", maxRetry = 3) {
+    cleanWorkspace()
+    writeLayout(
+      MillBuildLayout(
+        """|/MillMinimal/src/Main.scala
+           |package foo
+           |
+           |object Main {
+           |  def main(args: Array[String]): Unit = {
+           |     println("Hello java!")
+           |  }
+           |}
+           |/MillMinimal/test/src/Foo.scala
+           |// no test lense as debug is not supported
+           |class Foo extends munit.FunSuite {}
+           |""".stripMargin,
+        V.scala213,
+        V.millVersion,
+        includeMunit = true,
+      )
+    )
+
+    for {
+      _ <- server.initialize()
+      _ <- server.initialized()
+      _ <- server.executeCommand(ServerCommands.GenerateBspConfig)
+      _ <- server.didOpen("MillMinimal/src/Main.scala")
+      _ <- server.didSave("MillMinimal/src/Main.scala")(identity)
+      _ = assertNoDiagnostics()
+      _ <- assertCodeLenses(
+        "MillMinimal/src/Main.scala",
+        """|package foo
+           |
+           |<<run>>
+           |object Main {
+           |  def main(args: Array[String]): Unit = {
+           |     println("Hello java!")
+           |  }
+           |}""".stripMargin,
+      )
+      _ <- assertCodeLenses(
+        "MillMinimal/test/src/Foo.scala",
+        """|// no test lense as debug is not supported
+           |class Foo extends munit.FunSuite {}
+           |""".stripMargin,
+      )
+      lenses <- server.codeLenses("MillMinimal/src/Main.scala")
+      _ = assert(lenses.size > 0, "No lenses were generated!")
+      command = lenses.head.getCommand()
+      _ = assertEquals(runFromCommand(command), Some("Hello java!"))
+    } yield ()
+  }
+}

tests/unit/src/main/scala/tests/BaseCodeLensLspSuite.scala

@@ -2,10 +2,90 @@ package tests
 
 import scala.concurrent.Future
 
+import scala.meta.internal.builds.ShellRunner
+import scala.meta.internal.metals.MetalsEnrichments._
+
+import ch.epfl.scala.bsp4j.DebugSessionParams
+import com.google.gson.JsonObject
 import munit.Location
 import munit.TestOptions
+import org.eclipse.lsp4j.Command
+
+abstract class BaseCodeLensLspSuite(
+    name: String,
+    initializer: BuildServerInitializer = QuickBuildInitializer,
+) extends BaseLspSuite(name, initializer) {
+
+  protected def runFromCommand(cmd: Command): Option[String] = {
+    cmd.getArguments().asScala.toList match {
+      case (params: DebugSessionParams) :: _ =>
+        params.getData() match {
+          case obj: JsonObject =>
+            val cmd = obj
+              .get("shellCommand")
+              .getAsString()
+              // need to remove all escapes since ShellRunner does escaping itself
+              // for windows
+              .replace("\\\"", "")
+              .replace("\"\\", "\\")
+              // for linux
+              .replace("/\"", "/")
+              .replace("\"/", "/")
+              // when testing on Windows Program Files is problematic when splitting into list
+              .replace("Program Files", "ProgramFiles")
+              .replace("run shell command", "runshellcommand")
+              .split("\\s+")
+              .map(
+                // remove from escaped classpath
+                _.stripPrefix("\"")
+                  .stripSuffix("\"")
+                  // Add back spaces
+                  .replace("ProgramFiles", "Program Files")
+                  .replace("runshellcommand", "run shell command")
+              )
+            ShellRunner
+              .runSync(cmd.toList, workspace, redirectErrorOutput = false)
+              .map(_.trim())
+              .orElse {
+                scribe.error(
+                  "Couldn't run command specified in shellCommand."
+                )
+                scribe.error("The command run was:\n" + cmd.mkString(" "))
+                None
+              }
+          case _ => None
+        }
+
+      case _ => None
+    }
+  }
 
-abstract class BaseCodeLensLspSuite(name: String) extends BaseLspSuite(name) {
+  protected def testRunShellCommand(name: String): Unit =
+    test(name) {
+      cleanWorkspace()
+      for {
+        _ <- initialize(
+          s"""|/metals.json
+              |{
+              |  "a": {}
+              |}
+              |/a/src/main/scala/a/Main.scala
+              |package foo
+              |
+              |object Main {
+              |  def main(args: Array[String]): Unit = {
+              |     println("Hello java!")
+              |  }
+              |}
+              |""".stripMargin
+        )
+        _ <- server.didOpen("a/src/main/scala/a/Main.scala")
+        lenses <- server.codeLenses("a/src/main/scala/a/Main.scala")
+        _ = assert(lenses.size > 0, "No lenses were generated!")
+        command = lenses.head.getCommand()
+        _ = assertEquals(runFromCommand(command), Some("Hello java!"))
+      } yield ()
+    }
 
   def check(
       name: TestOptions,

tests/unit/src/main/scala/tests/BaseLspSuite.scala

@@ -112,7 +112,7 @@ abstract class BaseLspSuite(
 
   def test(
       testOpts: TestOptions,
-      withoutVirtualDocs: Boolean,
+      withoutVirtualDocs: Boolean = false,
       maxRetry: Int = 0,
   )(
       fn: => Future[Unit]