Let client supply code lenses it can run

As part of its initialization options the client can pass a
textDocument/codeLens object that lists the supported commands the
client can handle.

It is in the form of a dictionary where the key is the lens name
recognized by SourceKit-LSP and the value is the command as recognized
by the client.

```
initializationOptions: {
    "textDocument/codeLens": {
        supportedCommands: {
            "swift.run": "clientCommandName_Run",
            "swift.debug": "clientCommandName_Debug",
        }
    }
}
```

Author: plemarquand

Sources/LanguageServerProtocol/SupportTypes/ClientCapabilities.swift

@@ -465,6 +465,31 @@ public struct TextDocumentClientCapabilities: Hashable, Codable, Sendable {
     }
   }
 
+  public struct CodeLens: Hashable, Codable, Sendable {
+
+    /// Whether the client supports dynamic registration of this request.
+    public var dynamicRegistration: Bool?
+
+    public var supportedCommands: [String: String]
+
+    public init(dynamicRegistration: Bool? = nil, supportedCommands: [String: String] = [:]) {
+      self.dynamicRegistration = dynamicRegistration
+      self.supportedCommands = supportedCommands
+    }
+
+    public init(from decoder: any Decoder) throws {
+      let registration = try DynamicRegistrationCapability(from: decoder)
+      self = CodeLens(
+        dynamicRegistration: registration.dynamicRegistration
+      )
+    }
+
+    public func encode(to encoder: any Encoder) throws {
+      let registration = DynamicRegistrationCapability(dynamicRegistration: self.dynamicRegistration)
+      try registration.encode(to: encoder)
+    }
+  }
+
   /// Capabilities specific to `textDocument/rename`.
   public struct Rename: Hashable, Codable, Sendable {
 
@@ -666,7 +691,7 @@ public struct TextDocumentClientCapabilities: Hashable, Codable, Sendable {
 
   public var codeAction: CodeAction? = nil
 
-  public var codeLens: DynamicRegistrationCapability? = nil
+  public var codeLens: CodeLens? = nil
 
   public var documentLink: DynamicRegistrationCapability? = nil
 
@@ -715,7 +740,7 @@ public struct TextDocumentClientCapabilities: Hashable, Codable, Sendable {
     documentHighlight: DynamicRegistrationCapability? = nil,
     documentSymbol: DocumentSymbol? = nil,
     codeAction: CodeAction? = nil,
-    codeLens: DynamicRegistrationCapability? = nil,
+    codeLens: CodeLens? = nil,
     documentLink: DynamicRegistrationCapability? = nil,
     colorProvider: DynamicRegistrationCapability? = nil,
     formatting: DynamicRegistrationCapability? = nil,

Sources/SourceKitLSP/CapabilityRegistry.swift

@@ -69,10 +69,6 @@ public final actor CapabilityRegistry {
     clientCapabilities.textDocument?.diagnostic?.dynamicRegistration == true
   }
 
-  public var clientHasDynamicDocumentCodeLensRegistration: Bool {
-    clientCapabilities.textDocument?.codeLens?.dynamicRegistration == true
-  }
-
   public var clientHasDynamicExecuteCommandRegistration: Bool {
     clientCapabilities.workspace?.executeCommand?.dynamicRegistration == true
   }
@@ -87,6 +83,10 @@ public final actor CapabilityRegistry {
     clientCapabilities.textDocument?.publishDiagnostics?.codeDescriptionSupport == true
   }
 
+  public var supportedCodeLensCommands: [String: String] {
+    clientCapabilities.textDocument?.codeLens?.supportedCommands ?? [:]
+  }
+
   /// Since LSP 3.17.0, diagnostics can be reported through pull-based requests in addition to the existing push-based
   /// publish notifications.
   ///
@@ -322,29 +322,6 @@ public final actor CapabilityRegistry {
     )
   }
 
-  /// Dynamically register code lens capabilities,
-  /// if the client supports it.
-  public func registerCodeLensIfNeeded(
-    options: CodeLensOptions,
-    for languages: [Language],
-    server: SourceKitLSPServer
-  ) async {
-    guard clientHasDynamicDocumentCodeLensRegistration else { return }
-
-    await registerLanguageSpecificCapability(
-      options: CodeLensRegistrationOptions(
-        // Code lenses should only apply to saved files
-        documentSelector: DocumentSelector(for: languages, scheme: "file"),
-        codeLensOptions: options
-      ),
-      forMethod: CodeLensRequest.method,
-      languages: languages,
-      in: server,
-      registrationDict: codeLens,
-      setRegistrationDict: { codeLens[$0] = $1 }
-    )
-  }
-
   /// Dynamically register executeCommand with the given IDs if the client supports
   /// it and we haven't yet registered the given command IDs yet.
   public func registerExecuteCommandIfNeeded(

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -964,14 +964,30 @@ extension SourceKitLSPServer {
     // The below is a workaround for the vscode-swift extension since it cannot set client capabilities.
     // It passes "workspace/peekDocuments" through the `initializationOptions`.
     var clientCapabilities = req.capabilities
-    if case .dictionary(let initializationOptions) = req.initializationOptions,
-      let peekDocuments = initializationOptions["workspace/peekDocuments"]
-    {
-      if case .dictionary(var experimentalCapabilities) = clientCapabilities.experimental {
-        experimentalCapabilities["workspace/peekDocuments"] = peekDocuments
-        clientCapabilities.experimental = .dictionary(experimentalCapabilities)
-      } else {
-        clientCapabilities.experimental = .dictionary(["workspace/peekDocuments": peekDocuments])
+    if case .dictionary(let initializationOptions) = req.initializationOptions {
+      if let peekDocuments = initializationOptions["workspace/peekDocuments"] {
+        if case .dictionary(var experimentalCapabilities) = clientCapabilities.experimental {
+          experimentalCapabilities["workspace/peekDocuments"] = peekDocuments
+          clientCapabilities.experimental = .dictionary(experimentalCapabilities)
+        } else {
+          clientCapabilities.experimental = .dictionary(["workspace/peekDocuments": peekDocuments])
+        }
+      }
+
+      // The client announces what CodeLenses it supports, and the LSP will only return
+      // ones found in the supportedCommands dictionary.
+      if let codeLens = initializationOptions["textDocument/codeLens"],
+        case let .dictionary(codeLensConfig) = codeLens,
+        case let .dictionary(supportedCommands) = codeLensConfig["supportedCommands"]
+      {
+        let commandMap = supportedCommands.compactMapValues({
+          if case let .string(val) = $0 {
+            return val
+          }
+          return nil
+        })
+
+        clientCapabilities.textDocument?.codeLens?.supportedCommands = commandMap
       }
     }
 
@@ -1103,7 +1119,7 @@ extension SourceKitLSPServer {
           supportsCodeActions: true
         )
       ),
-      codeLensProvider: CodeLensOptions(resolveProvider: false),
+      codeLensProvider: CodeLensOptions(),
       documentFormattingProvider: .value(DocumentFormattingOptions(workDoneProgress: false)),
       renameProvider: .value(RenameOptions(prepareProvider: true)),
       colorProvider: .bool(true),
@@ -1164,9 +1180,6 @@ extension SourceKitLSPServer {
     if let diagnosticOptions = server.diagnosticProvider {
       await registry.registerDiagnosticIfNeeded(options: diagnosticOptions, for: languages, server: self)
     }
-    if let codeLensOptions = server.codeLensProvider {
-      await registry.registerCodeLensIfNeeded(options: codeLensOptions, for: languages, server: self)
-    }
     if let commandOptions = server.executeCommandProvider {
       await registry.registerExecuteCommandIfNeeded(commands: commandOptions.commands, server: self)
     }

Sources/SourceKitLSP/Swift/SwiftCodeLensScanner.swift

@@ -16,29 +16,34 @@ import SwiftSyntax
 /// Scans a source file for classes or structs annotated with `@main` and returns a code lens for them.
 final class SwiftCodeLensScanner: SyntaxVisitor {
   /// The document snapshot of the syntax tree that is being walked.
-  private var snapshot: DocumentSnapshot
+  private let snapshot: DocumentSnapshot
 
   /// The collection of CodeLenses found in the document.
   private var result: [CodeLens] = []
 
-  private init(snapshot: DocumentSnapshot) {
+  /// The map of supported commands and their client side command names
+  private let supportedCommands: [String: String]
+
+  private init(snapshot: DocumentSnapshot, supportedCommands: [String: String]) {
     self.snapshot = snapshot
+    self.supportedCommands = supportedCommands
     super.init(viewMode: .fixedUp)
   }
 
   /// Public entry point. Scans the syntax tree of the given snapshot for an `@main` annotation
   /// and returns CodeLens's with Commands to run/debug the application.
   public static func findCodeLenses(
     in snapshot: DocumentSnapshot,
-    syntaxTreeManager: SyntaxTreeManager
+    syntaxTreeManager: SyntaxTreeManager,
+    supportedCommands: [String: String]
   ) async -> [CodeLens] {
-    guard snapshot.text.contains("@main") else {
+    guard snapshot.text.contains("@main") && !supportedCommands.isEmpty else {
       // This is intended to filter out files that obviously do not contain an entry point.
       return []
     }
 
     let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot)
-    let visitor = SwiftCodeLensScanner(snapshot: snapshot)
+    let visitor = SwiftCodeLensScanner(snapshot: snapshot, supportedCommands: supportedCommands)
     visitor.walk(syntaxTree)
     return visitor.result
   }
@@ -57,21 +62,25 @@ final class SwiftCodeLensScanner: SyntaxVisitor {
     if attribute.trimmedDescription == "@main" {
       let range = self.snapshot.absolutePositionRange(of: attribute.trimmedRange)
 
-      // Return commands for running/debugging the executable.
-      // These command names must be recognized by the client and so should not be chosen arbitrarily.
-      self.result.append(
-        CodeLens(
-          range: range,
-          command: Command(title: "Run", command: "swift.run", arguments: nil)
+      if let runCommand = supportedCommands["swift.run"] {
+        // Return commands for running/debugging the executable.
+        // These command names must be recognized by the client and so should not be chosen arbitrarily.
+        self.result.append(
+          CodeLens(
+            range: range,
+            command: Command(title: "Run", command: runCommand, arguments: nil)
+          )
         )
-      )
+      }
 
-      self.result.append(
-        CodeLens(
-          range: range,
-          command: Command(title: "Debug", command: "swift.debug", arguments: nil)
+      if let debugCommand = supportedCommands["swift.debug"] {
+        self.result.append(
+          CodeLens(
+            range: range,
+            command: Command(title: "Debug", command: debugCommand, arguments: nil)
+          )
         )
-      )
+      }
     }
   }
 }

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -317,6 +317,7 @@ extension SwiftLanguageService {
             supportsCodeActions: true
           )
         ),
+        codeLensProvider: CodeLensOptions(),
         colorProvider: .bool(true),
         foldingRangeProvider: .bool(true),
         executeCommandProvider: ExecuteCommandOptions(
@@ -923,7 +924,11 @@ extension SwiftLanguageService {
 
   public func codeLens(_ req: CodeLensRequest) async throws -> [CodeLens] {
     let snapshot = try documentManager.latestSnapshot(req.textDocument.uri)
-    return await SwiftCodeLensScanner.findCodeLenses(in: snapshot, syntaxTreeManager: self.syntaxTreeManager)
+    return await SwiftCodeLensScanner.findCodeLenses(
+      in: snapshot,
+      syntaxTreeManager: self.syntaxTreeManager,
+      supportedCommands: self.capabilityRegistry.supportedCodeLensCommands
+    )
   }
 
   public func documentDiagnostic(_ req: DocumentDiagnosticsRequest) async throws -> DocumentDiagnosticReport {

Tests/SourceKitLSPTests/CodeLensTests.swift

@@ -17,15 +17,44 @@ import XCTest
 
 final class CodeLensTests: XCTestCase {
   func testNoLenses() async throws {
+    var codeLensCapabilities = TextDocumentClientCapabilities.CodeLens()
+    codeLensCapabilities.supportedCommands = [
+      "swift.run": "swift.run",
+      "swift.debug": "swift.debug"
+    ]
+    let capabilities = ClientCapabilities(textDocument: TextDocumentClientCapabilities(codeLens: codeLensCapabilities))
+
     let project = try await SwiftPMTestProject(
       files: [
         "Test.swift": """
         struct MyApp {
           public static func main() {}
         }
         """
+      ],
+      capabilities: capabilities
+    )
+    let (uri, _) = try project.openDocument("Test.swift")
+
+    let response = try await project.testClient.send(
+      CodeLensRequest(textDocument: TextDocumentIdentifier(uri))
+    )
+
+    XCTAssertEqual(response, [])
+  }
+
+  func testNoClientCodeLenses() async throws {
+    let project = try await SwiftPMTestProject(
+      files: [
+        "Test.swift": """
+        @main
+        struct MyApp {
+          public static func main() {}
+        }
+        """
       ]
     )
+
     let (uri, _) = try project.openDocument("Test.swift")
 
     let response = try await project.testClient.send(
@@ -36,6 +65,13 @@ final class CodeLensTests: XCTestCase {
   }
 
   func testSuccessfulCodeLensRequest() async throws {
+    var codeLensCapabilities = TextDocumentClientCapabilities.CodeLens()
+    codeLensCapabilities.supportedCommands = [
+      "swift.run": "swift.run",
+      "swift.debug": "swift.debug"
+    ]
+    let capabilities = ClientCapabilities(textDocument: TextDocumentClientCapabilities(codeLens: codeLensCapabilities))
+
     let project = try await SwiftPMTestProject(
       files: [
         "Test.swift": """
@@ -44,7 +80,8 @@ final class CodeLensTests: XCTestCase {
           public static func main() {}
         }
         """
-      ]
+      ],
+      capabilities: capabilities
     )
 
     let (uri, positions) = try project.openDocument("Test.swift")