Don’t catch CancellationError for document diagnostics

We were catching errors during diagnostic generation because VS Code will not request diagnostics again if SourceKit-LSP returns an error to any diagnostic request. We were, however, a little over-eager when doing this and also caught cancellation errors, which means that if the user made multiple edits in quick succession (which would cancel the diagnostics request from the first edit), we would return empty diagnostics, clearing the displayed diagnostics.

Author: ahoppen

Sources/SKTestSupport/TestSourceKitLSPClient.swift

@@ -206,10 +206,16 @@ public final class TestSourceKitLSPClient: MessageHandler {
   ///
   /// This version of the `send` function should only be used if some action needs to be performed after the request is
   /// sent but before it returns a result.
-  public func send<R: RequestType>(_ request: R, completionHandler: @escaping (LSPResult<R.Response>) -> Void) {
-    server.handle(request, id: .number(Int(nextRequestID.fetchAndIncrement()))) { result in
+  @discardableResult
+  public func send<R: RequestType>(
+    _ request: R,
+    completionHandler: @escaping (LSPResult<R.Response>) -> Void
+  ) -> RequestID {
+    let requestID = RequestID.number(Int(nextRequestID.fetchAndIncrement()))
+    server.handle(request, id: requestID) { result in
       completionHandler(result)
     }
+    return requestID
   }
 
   /// Send the notification to `server`.

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -888,6 +888,8 @@ extension SwiftLanguageService {
         buildSettings: buildSettings
       )
       return .full(diagnosticReport)
+    } catch let error as CancellationError {
+      throw error
     } catch {
       // VS Code does not request diagnostics again for a document if the diagnostics request failed.
       // Since sourcekit-lsp usually recovers from failures (e.g. after sourcekitd crashes), this is undesirable.

Tests/SourceKitLSPTests/PullDiagnosticsTests.swift

@@ -309,4 +309,32 @@ final class PullDiagnosticsTests: XCTestCase {
     diagnosticRequestSent.value = true
     try await fulfillmentOfOrThrow([receivedDiagnostics])
   }
+
+  func testDontReturnEmptyDiagnosticsIfDiagnosticRequestIsCancelled() async throws {
+    let diagnosticRequestCancelled = self.expectation(description: "diagnostic request cancelled")
+    var serverOptions = SourceKitLSPServer.Options.testDefault
+    serverOptions.indexTestHooks.preparationTaskDidStart = { _ in
+      await self.fulfillment(of: [diagnosticRequestCancelled], timeout: defaultTimeout)
+    }
+    let project = try await SwiftPMTestProject(
+      files: [
+        "Lib.swift": "let x: String = 1"
+      ],
+      serverOptions: serverOptions,
+      enableBackgroundIndexing: true,
+      pollIndex: false
+    )
+    let (uri, _) = try project.openDocument("Lib.swift")
+
+    let diagnosticResponseReceived = self.expectation(description: "Received diagnostic response")
+    let requestID = project.testClient.send(
+      DocumentDiagnosticsRequest(textDocument: TextDocumentIdentifier(uri))
+    ) { result in
+      XCTAssertEqual(result.failure?.code, .cancelled)
+      diagnosticResponseReceived.fulfill()
+    }
+    project.testClient.send(CancelRequestNotification(id: requestID))
+    diagnosticRequestCancelled.fulfill()
+    try await fulfillmentOfOrThrow([diagnosticResponseReceived])
+  }
 }