Don’t catch CancellationError for document diagnostics

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])
+  }
 }