Only provide semantic tokens if capability is registered

Author: fwcd

Sources/SourceKitLSP/CapabilityRegistry.swift

@@ -142,6 +142,11 @@ public final class CapabilityRegistry {
     }
   }
 
+  /// Checks whether a registration for semantic tokens for the given languages exists.
+  public func hasSemanticTokensRegistration(for languages: [Language]) -> Bool {
+    registration(for: languages, in: semanticTokens) != nil
+  }
+
   private func documentSelector(for langauges: [Language]) -> DocumentSelector {
     return DocumentSelector(langauges.map { DocumentFilter(language: $0.rawValue) })
   }

Sources/SourceKitLSP/SourceKitServer.swift

@@ -855,28 +855,49 @@ extension SourceKitServer {
     languageService.documentColor(req)
   }
 
+  private func supportsSemanticTokens(in uri: DocumentURI) -> Bool {
+    guard let workspace = workspace,
+          let snapshot = workspace.documentManager.latestSnapshot(uri) else {
+      return false
+    }
+    let language = snapshot.document.language
+    return workspace.capabilityRegistry.hasSemanticTokensRegistration(for: [language])
+  }
+
   func documentSemanticTokens(
     _ req: Request<DocumentSemanticTokensRequest>,
     workspace: Workspace,
     languageService: ToolchainLanguageServer
   ) {
-    languageService.documentSemanticTokens(req)
+    if supportsSemanticTokens(in: req.params.textDocument.uri) {
+      languageService.documentSemanticTokens(req)
+    } else {
+      req.reply(DocumentSemanticTokensResponse(data: []))
+    }
   }
 
   func documentSemanticTokensDelta(
     _ req: Request<DocumentSemanticTokensDeltaRequest>,
     workspace: Workspace,
     languageService: ToolchainLanguageServer
   ) {
-    languageService.documentSemanticTokensDelta(req)
+    if supportsSemanticTokens(in: req.params.textDocument.uri) {
+      languageService.documentSemanticTokensDelta(req)
+    } else {
+      req.reply(.tokens(.init(data: [])))
+    }
   }
 
   func documentSemanticTokensRange(
     _ req: Request<DocumentSemanticTokensRangeRequest>,
     workspace: Workspace,
     languageService: ToolchainLanguageServer
   ) {
-    languageService.documentSemanticTokensRange(req)
+    if supportsSemanticTokens(in: req.params.textDocument.uri) {
+      languageService.documentSemanticTokensRange(req)
+    } else {
+      req.reply(DocumentSemanticTokensResponse(data: []))
+    }
   }
 
   func colorPresentation(

Sources/SourceKitLSP/Swift/SyntaxHighlightingToken.swift

@@ -104,7 +104,8 @@ public struct SyntaxHighlightingToken: Hashable {
     case `operator`
 
     /// The name of the token type used by LSP.
-    var lspName: String {
+    /// **Public for testing.**
+    public var lspName: String {
       switch self {
       case .namespace: return "namespace"
       case .type: return "type"
@@ -168,6 +169,7 @@ public struct SyntaxHighlightingToken: Hashable {
     /// The name of the modifier used by LSP, if this
     /// is a single modifier. Note that every modifier
     /// in `allCases` must have an associated `lspName`.
+    /// **Public for testing.**
     public var lspName: String? {
       switch self {
       case .declaration: return "declaration"

Tests/SourceKitLSPTests/SemanticTokensTests.swift

@@ -44,7 +44,18 @@ final class SemanticTokensTests: XCTestCase {
             refreshSupport: true
           )
         ),
-        textDocument: nil
+        textDocument: .init(
+          semanticTokens: .init(
+            dynamicRegistration: true,
+            requests: .init(
+              range: .bool(true),
+              full: .bool(true)
+            ),
+            tokenTypes: Token.Kind.allCases.map(\.lspName),
+            tokenModifiers: Token.Modifiers.allCases.map { $0.lspName! },
+            formats: [.relative]
+          )
+        )
       ),
       trace: .off,
       workspaceFolders: nil
@@ -54,12 +65,24 @@ final class SemanticTokensTests: XCTestCase {
   private func performSemanticTokensRequest(text: String, range: Range<Position>? = nil) -> [Token] {
     let url = URL(fileURLWithPath: "/\(#function)/a.swift")
 
-    // We wait for the first refresh request to make sure that the semantic tokens are ready
+    // We will wait for the server to dynamically register semantic tokens
+
+    let registerCapabilityExpectation = expectation(description: "performSemanticTokensRequest - register semantic tokens capability")
+    sk.appendOneShotRequestHandler { (req: Request<RegisterCapabilityRequest>) in
+      let registrations = req.params.registrations
+      XCTAssert(registrations.contains { reg in
+        reg.method == SemanticTokensRegistrationOptions.method
+      })
+      req.reply(VoidResponse())
+      registerCapabilityExpectation.fulfill()
+    }
+
+    // We will wait for the first refresh request to make sure that the semantic tokens are ready
 
     let refreshExpectation = expectation(description: "performSemanticTokensRequest - refresh received")
-    sk.handleNextRequest { (req: Request<WorkspaceSemanticTokensRefreshRequest>) in
-      refreshExpectation.fulfill()
+    sk.appendOneShotRequestHandler { (req: Request<WorkspaceSemanticTokensRefreshRequest>) in
       req.reply(VoidResponse())
+      refreshExpectation.fulfill()
     }
 
     sk.send(DidOpenTextDocumentNotification(textDocument: TextDocumentItem(
@@ -69,7 +92,7 @@ final class SemanticTokensTests: XCTestCase {
       text: text
     )))
 
-    wait(for: [refreshExpectation], timeout: 15)
+    wait(for: [registerCapabilityExpectation, refreshExpectation], timeout: 15)
 
     let textDocument = TextDocumentIdentifier(url)
     let response: DocumentSemanticTokensResponse!