Do not block initialization of the build server when build server is unresponsive in returning the list of test files

Sources/BuildSystemIntegration/BuildSystemHooks.swift

@@ -31,12 +31,6 @@ package struct SwiftPMTestHooks: Sendable {
   }
 }
 
-/// When running SourceKit-LSP in-process, allows the creator of `SourceKitLSPServer` to create the build system instead
-/// of SourceKit-LSP creating build systems as needed.
-package protocol BuildSystemInjector: Sendable {
-  func createBuildSystem(projectRoot: URL, connectionToSourceKitLSP: any Connection) async -> BuiltInBuildSystem
-}
-
 package struct BuildSystemHooks: Sendable {
   package var swiftPMTestHooks: SwiftPMTestHooks
 
@@ -45,15 +39,20 @@ package struct BuildSystemHooks: Sendable {
   /// This allows requests to be artificially delayed.
   package var preHandleRequest: (@Sendable (any RequestType) async -> Void)?
 
-  package var buildSystemInjector: BuildSystemInjector?
+  /// When running SourceKit-LSP in-process, allows the creator of `SourceKitLSPServer` to create a message handler that
+  /// handles BSP requests instead of SourceKit-LSP creating build systems as needed.
+  package var injectBuildServer:
+    (@Sendable (_ projectRoot: URL, _ connectionToSourceKitLSP: any Connection) async -> any Connection)?
 
   package init(
     swiftPMTestHooks: SwiftPMTestHooks = SwiftPMTestHooks(),
     preHandleRequest: (@Sendable (any RequestType) async -> Void)? = nil,
-    buildSystemInjector: BuildSystemInjector? = nil
+    injectBuildServer: (
+      @Sendable (_ projectRoot: URL, _ connectionToSourceKitLSP: any Connection) async -> any Connection
+    )? = nil
   ) {
     self.swiftPMTestHooks = swiftPMTestHooks
     self.preHandleRequest = preHandleRequest
-    self.buildSystemInjector = buildSystemInjector
+    self.injectBuildServer = injectBuildServer
   }
 }

Sources/BuildSystemIntegration/BuildSystemManager.swift

@@ -119,6 +119,8 @@ fileprivate extension InitializeBuildResponse {
 private enum BuildSystemAdapter {
   case builtIn(BuiltInBuildSystemAdapter, connectionToBuildSystem: any Connection)
   case external(ExternalBuildSystemAdapter)
+  /// A message handler that was created by `injectBuildServer` and will handle all BSP messages.
+  case injected(any Connection)
 
   /// Send a notification to the build server.
   func send(_ notification: some NotificationType) async {
@@ -127,6 +129,8 @@ private enum BuildSystemAdapter {
       connectionToBuildSystem.send(notification)
     case .external(let external):
       await external.send(notification)
+    case .injected(let connection):
+      connection.send(notification)
     }
   }
 
@@ -137,6 +141,33 @@ private enum BuildSystemAdapter {
       return try await connectionToBuildSystem.send(request)
     case .external(let external):
       return try await external.send(request)
+    case .injected(let messageHandler):
+      // After we sent the request, the ID of the request.
+      // When we send a `CancelRequestNotification` this is reset to `nil` so that we don't send another cancellation
+      // notification.
+      let requestID = ThreadSafeBox<RequestID?>(initialValue: nil)
+
+      return try await withTaskCancellationHandler {
+        return try await withCheckedThrowingContinuation { continuation in
+          if Task.isCancelled {
+            return continuation.resume(throwing: CancellationError())
+          }
+          requestID.value = messageHandler.send(request) { response in
+            continuation.resume(with: response)
+          }
+          if Task.isCancelled {
+            // The task might have been cancelled after we checked `Task.isCancelled` above but before `requestID.value`
+            // is set, we won't send a `CancelRequestNotification` from the `onCancel` handler. Send it from here.
+            if let requestID = requestID.takeValue() {
+              messageHandler.send(CancelRequestNotification(id: requestID))
+            }
+          }
+        }
+      } onCancel: {
+        if let requestID = requestID.takeValue() {
+          messageHandler.send(CancelRequestNotification(id: requestID))
+        }
+      }
     }
   }
 }
@@ -234,12 +265,13 @@ private extension BuildSystemSpec {
       return nil
       #endif
     case .injected(let injector):
-      return await createBuiltInBuildSystemAdapter(
-        messagesToSourceKitLSPHandler: messagesToSourceKitLSPHandler,
-        buildSystemHooks: buildSystemHooks
-      ) { connectionToSourceKitLSP in
-        await injector.createBuildSystem(projectRoot: projectRoot, connectionToSourceKitLSP: connectionToSourceKitLSP)
-      }
+      let connectionToSourceKitLSP = LocalConnection(
+        receiverName: "BuildSystemManager for \(projectRoot.lastPathComponent)"
+      )
+      connectionToSourceKitLSP.start(handler: messagesToSourceKitLSPHandler)
+      return .injected(
+        await injector(projectRoot, connectionToSourceKitLSP)
+      )
     }
   }
 }

Sources/BuildSystemIntegration/BuiltInBuildSystemAdapter.swift

@@ -11,7 +11,6 @@
 //===----------------------------------------------------------------------===//
 
 import BuildServerProtocol
-import LanguageServerProtocol
 import LanguageServerProtocolExtensions
 import SKLogging
 import SKOptions
@@ -20,8 +19,10 @@ import ToolchainRegistry
 
 #if compiler(>=6)
 package import Foundation
+package import LanguageServerProtocol
 #else
 import Foundation
+import LanguageServerProtocol
 #endif
 
 /// The details necessary to create a `BuildSystemAdapter`.
@@ -31,7 +32,9 @@ package struct BuildSystemSpec {
     case jsonCompilationDatabase
     case fixedCompilationDatabase
     case swiftPM
-    case injected(BuildSystemInjector)
+    case injected(
+      @Sendable (_ projectRoot: URL, _ connectionToSourceKitLSP: any Connection) async -> any Connection
+    )
   }
 
   package var kind: Kind

Sources/BuildSystemIntegration/DetermineBuildSystem.swift

@@ -72,9 +72,9 @@ package func determineBuildSystem(
   options: SourceKitLSPOptions,
   hooks: BuildSystemHooks
 ) -> BuildSystemSpec? {
-  if let buildSystemInjector = hooks.buildSystemInjector {
+  if let injectBuildServer = hooks.injectBuildServer {
     return BuildSystemSpec(
-      kind: .injected(buildSystemInjector),
+      kind: .injected(injectBuildServer),
       projectRoot: workspaceFolder.arbitrarySchemeURL,
       configPath: workspaceFolder.arbitrarySchemeURL
     )

Sources/BuildSystemIntegration/TestBuildSystem.swift

@@ -10,28 +10,23 @@
 //
 //===----------------------------------------------------------------------===//
 
+import Foundation
+import LanguageServerProtocolExtensions
+import SKLogging
+import SKOptions
+import ToolchainRegistry
+
 #if compiler(>=6)
 package import BuildServerProtocol
-package import Foundation
 package import LanguageServerProtocol
-import SKOptions
-import ToolchainRegistry
 #else
 import BuildServerProtocol
-import Foundation
 import LanguageServerProtocol
-import SKOptions
-import ToolchainRegistry
 #endif
 
 /// Build system to be used for testing BuildSystem and BuildSystemDelegate functionality with SourceKitLSPServer
 /// and other components.
-package actor TestBuildSystem: BuiltInBuildSystem {
-  package let fileWatchers: [FileSystemWatcher] = []
-
-  package let indexStorePath: URL? = nil
-  package let indexDatabasePath: URL? = nil
-
+package actor TestBuildSystem: MessageHandler {
   private let connectionToSourceKitLSP: any Connection
 
   /// Build settings by file.
@@ -48,7 +43,83 @@ package actor TestBuildSystem: BuiltInBuildSystem {
     self.connectionToSourceKitLSP = connectionToSourceKitLSP
   }
 
-  package func buildTargets(request: WorkspaceBuildTargetsRequest) async throws -> WorkspaceBuildTargetsResponse {
+  package nonisolated func handle(_ notification: some NotificationType) {
+    do {
+      switch notification {
+      case let notification as CancelRequestNotification:
+        try self.cancelRequest(notification)
+      case let notification as OnBuildExitNotification:
+        try self.onBuildExit(notification)
+      case let notification as OnBuildInitializedNotification:
+        try self.onBuildInitialized(notification)
+      case let notification as OnWatchedFilesDidChangeNotification:
+        try self.onWatchedFilesDidChange(notification)
+      default:
+        throw ResponseError.methodNotFound(type(of: notification).method)
+      }
+    } catch {
+      logger.error("Error while handling BSP notification")
+    }
+  }
+
+  package nonisolated func handle<Request: RequestType>(
+    _ request: Request,
+    id: RequestID,
+    reply: @Sendable @escaping (LSPResult<Request.Response>) -> Void
+  ) {
+    func handle<R: RequestType>(_ request: R, using handler: @Sendable @escaping (R) async throws -> R.Response) {
+      Task {
+        do {
+          reply(.success(try await handler(request) as! Request.Response))
+        } catch {
+          reply(.failure(ResponseError(error)))
+        }
+      }
+    }
+
+    switch request {
+    case let request as BuildShutdownRequest:
+      handle(request, using: self.buildShutdown(_:))
+    case let request as BuildTargetSourcesRequest:
+      handle(request, using: self.buildTargetSourcesRequest)
+    case let request as InitializeBuildRequest:
+      handle(request, using: self.initializeBuildRequest)
+    case let request as TextDocumentSourceKitOptionsRequest:
+      handle(request, using: self.textDocumentSourceKitOptionsRequest)
+    case let request as WorkspaceBuildTargetsRequest:
+      handle(request, using: self.workspaceBuildTargetsRequest)
+    case let request as WorkspaceWaitForBuildSystemUpdatesRequest:
+      handle(request, using: self.workspaceWaitForBuildSystemUpdatesRequest)
+    default:
+      reply(.failure(ResponseError.methodNotFound(type(of: request).method)))
+    }
+  }
+
+  func initializeBuildRequest(_ request: InitializeBuildRequest) async throws -> InitializeBuildResponse {
+    return InitializeBuildResponse(
+      displayName: "TestBuildSystem",
+      version: "",
+      bspVersion: "2.2.0",
+      capabilities: BuildServerCapabilities(),
+      data: nil
+    )
+  }
+
+  nonisolated func onBuildInitialized(_ notification: OnBuildInitializedNotification) throws {
+    // Nothing to do
+  }
+
+  func buildShutdown(_ request: BuildShutdownRequest) async throws -> VoidResponse {
+    return VoidResponse()
+  }
+
+  nonisolated func onBuildExit(_ notification: OnBuildExitNotification) throws {
+    // Nothing to do
+  }
+
+  func workspaceBuildTargetsRequest(
+    _ request: WorkspaceBuildTargetsRequest
+  ) async throws -> WorkspaceBuildTargetsResponse {
     return WorkspaceBuildTargetsResponse(targets: [
       BuildTarget(
         id: .dummy,
@@ -62,7 +133,7 @@ package actor TestBuildSystem: BuiltInBuildSystem {
     ])
   }
 
-  package func buildTargetSources(request: BuildTargetSourcesRequest) async throws -> BuildTargetSourcesResponse {
+  func buildTargetSourcesRequest(_ request: BuildTargetSourcesRequest) async throws -> BuildTargetSourcesResponse {
     return BuildTargetSourcesResponse(items: [
       SourcesItem(
         target: .dummy,
@@ -71,19 +142,25 @@ package actor TestBuildSystem: BuiltInBuildSystem {
     ])
   }
 
-  package func didChangeWatchedFiles(notification: OnWatchedFilesDidChangeNotification) async {}
-
-  package func prepare(request: BuildTargetPrepareRequest) async throws -> VoidResponse {
-    throw PrepareNotSupportedError()
-  }
-
-  package func sourceKitOptions(
-    request: TextDocumentSourceKitOptionsRequest
+  func textDocumentSourceKitOptionsRequest(
+    _ request: TextDocumentSourceKitOptionsRequest
   ) async throws -> TextDocumentSourceKitOptionsResponse? {
     return buildSettingsByFile[request.textDocument.uri]
   }
 
   package func waitForBuildSystemUpdates(request: WorkspaceWaitForBuildSystemUpdatesRequest) async -> VoidResponse {
     return VoidResponse()
   }
+
+  nonisolated func onWatchedFilesDidChange(_ notification: OnWatchedFilesDidChangeNotification) throws {
+    // Not watching any files
+  }
+
+  func workspaceWaitForBuildSystemUpdatesRequest(
+    _ request: WorkspaceWaitForBuildSystemUpdatesRequest
+  ) async throws -> VoidResponse {
+    return VoidResponse()
+  }
+
+  nonisolated func cancelRequest(_ notification: CancelRequestNotification) throws {}
 }

Sources/SKTestSupport/TestSourceKitLSPClient.swift

@@ -204,17 +204,20 @@ package final class TestSourceKitLSPClient: MessageHandler, Sendable {
     }
     preInitialization?(self)
     if initialize {
-      _ = try await self.send(
-        InitializeRequest(
-          processId: nil,
-          rootPath: nil,
-          rootURI: nil,
-          initializationOptions: initializationOptions,
-          capabilities: capabilities,
-          trace: .off,
-          workspaceFolders: workspaceFolders
+      let capabilities = capabilities
+      try await withTimeout(.seconds(defaultTimeout)) {
+        _ = try await self.send(
+          InitializeRequest(
+            processId: nil,
+            rootPath: nil,
+            rootURI: nil,
+            initializationOptions: initializationOptions,
+            capabilities: capabilities,
+            trace: .off,
+            workspaceFolders: workspaceFolders
+          )
         )
-      )
+      }
     }
   }