Add session API for registering additional toolchains

Sources/SWBBuildService/Messages.swift

@@ -302,6 +302,13 @@ private struct SetSessionUserPreferencesMsg: MessageHandler {
     }
 }
 
+private struct RegisterToolchainHandler: MessageHandler {
+    func handle(request: Request, message: RegisterToolchainRequest) async throws -> StringResponse {
+        let session = try request.session(for: message)
+        return StringResponse(try await session.core.registerToolchain(at: message.path))
+    }
+}
+
 /// Start a PIF transfer from the client.
 ///
 /// This will establish a workspace context in the relevant session by exchanging a PIF from the client to the service incrementally, only transferring subobjects as necessary.
@@ -1542,6 +1549,7 @@ public struct ServiceSessionMessageHandlers: ServiceExtension {
         service.registerMessageHandler(SetSessionSystemInfoMsg.self)
         service.registerMessageHandler(SetSessionUserInfoMsg.self)
         service.registerMessageHandler(SetSessionUserPreferencesMsg.self)
+        service.registerMessageHandler(RegisterToolchainHandler.self)
         service.registerMessageHandler(DeveloperPathHandler.self)
     }
 }

Sources/SWBCore/Core.swift

@@ -533,6 +533,10 @@ public final class Core: Sendable {
         specRegistry.freeze()
     }
 
+    public func registerToolchain(at toolchainPath: Path) async throws -> String {
+        return try await toolchainRegistry.registerToolchain(at: toolchainPath, operatingSystem: hostOperatingSystem, delegate: registryDelegate, diagnoseAlreadyRegisteredToolchain: false, aliases: [])
+    }
+
     /// Dump information on the registered platforms.
     public func getPlatformsDump() -> String {
         var result = ""
@@ -615,7 +619,7 @@ public final class Core: Sendable {
     /// Dump information on the registered toolchains.
     public func getToolchainsDump() async -> String {
         var result = ""
-        for (_,toolchain) in toolchainRegistry.toolchainsByIdentifier.sorted(byKey: <) {
+        for toolchain in toolchainRegistry.toolchains.sorted(by: \.identifier) {
             result += "\(toolchain)\n"
         }
         return result

Sources/SWBCore/ToolchainRegistry.swift

@@ -11,7 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 public import SWBUtil
-
+import Synchronization
 import struct Foundation.CharacterSet
 import Foundation
 import SWBMacro
@@ -428,6 +428,15 @@ extension Toolchain {
 
 /// The ToolchainRegistry manages the set of registered toolchains.
 public final class ToolchainRegistry: @unchecked Sendable {
+    enum Error: Swift.Error, CustomStringConvertible {
+        case toolchainAlreadyRegistered(String, Path)
+
+        var description: String {
+            switch self {
+            case .toolchainAlreadyRegistered(let identifier, let path): "toolchain '\(identifier)' already registered from \(path.str)"
+            }
+        }
+    }
     @_spi(Testing) public struct SearchPath: Sendable {
         public var path: Path
         public var strict: Bool
@@ -443,11 +452,14 @@ public final class ToolchainRegistry: @unchecked Sendable {
     let fs: any FSProxy
     let hostOperatingSystem: OperatingSystem
 
-    /// The map of toolchains by identifier.
-    @_spi(Testing) public private(set) var toolchainsByIdentifier = Dictionary<String, Toolchain>()
+    struct State {
+        /// The map of toolchains by identifier.
+        @_spi(Testing) public fileprivate(set) var toolchainsByIdentifier = Dictionary<String, Toolchain>()
 
-    /// Lower-cased alias -> toolchain (alias lookup is case-insensitive)
-    @_spi(Testing) public private(set) var toolchainsByAlias = Dictionary<String, Toolchain>()
+        /// Lower-cased alias -> toolchain (alias lookup is case-insensitive)
+        @_spi(Testing) public fileprivate(set) var toolchainsByAlias = Dictionary<String, Toolchain>()
+    }
+    private let state: SWBMutex<State> = .init(State())
 
     public static let defaultToolchainIdentifier: String = "com.apple.dt.toolchain.XcodeDefault"
 
@@ -503,40 +515,53 @@ public final class ToolchainRegistry: @unchecked Sendable {
             guard toolchainPath.basenameWithoutSuffix != "swift-latest" else { continue }
 
             do {
-                let toolchain = try await Toolchain(path: toolchainPath, operatingSystem: operatingSystem, aliases: aliases, fs: fs, pluginManager: delegate.pluginManager, platformRegistry: delegate.platformRegistry)
-                try register(toolchain)
+                _ = try await registerToolchain(at: toolchainPath, operatingSystem: operatingSystem, delegate: delegate, diagnoseAlreadyRegisteredToolchain: true, aliases: aliases)
             } catch let err {
                 delegate.issue(strict: strict, toolchainPath, "failed to load toolchain: \(err)")
             }
         }
     }
 
-    private func register(_ toolchain: Toolchain) throws {
-        if let duplicateToolchain = toolchainsByIdentifier[toolchain.identifier] {
-            throw StubError.error("toolchain '\(toolchain.identifier)' already registered from \(duplicateToolchain.path.str)")
+    func registerToolchain(at toolchainPath: Path, operatingSystem: OperatingSystem, delegate: any ToolchainRegistryDelegate, diagnoseAlreadyRegisteredToolchain: Bool, aliases: Set<String>) async throws -> String {
+        do {
+            let toolchain = try await Toolchain(path: toolchainPath, operatingSystem: operatingSystem, aliases: aliases, fs: fs, pluginManager: delegate.pluginManager, platformRegistry: delegate.platformRegistry)
+            try register(toolchain)
+            return toolchain.identifier
+        } catch Error.toolchainAlreadyRegistered(let identifier, _) where !diagnoseAlreadyRegisteredToolchain {
+            return identifier
         }
-        toolchainsByIdentifier[toolchain.identifier] = toolchain
-
-        for alias in toolchain.aliases {
-            guard !alias.isEmpty else { continue }
-            assert(alias.lowercased() == alias)
+    }
 
-            // When two toolchains have conflicting aliases, the highest-versioned toolchain wins (regardless of identifier)
-            if let existingToolchain = toolchainsByAlias[alias], existingToolchain.version >= toolchain.version {
-                continue
+    private func register(_ toolchain: Toolchain) throws {
+        try state.withLock { state in
+            if let duplicateToolchain = state.toolchainsByIdentifier[toolchain.identifier] {
+                throw Error.toolchainAlreadyRegistered(toolchain.identifier, duplicateToolchain.path)
             }
+            state.toolchainsByIdentifier[toolchain.identifier] = toolchain
+
+            for alias in toolchain.aliases {
+                guard !alias.isEmpty else { continue }
+                assert(alias.lowercased() == alias)
+
+                // When two toolchains have conflicting aliases, the highest-versioned toolchain wins (regardless of identifier)
+                if let existingToolchain = state.toolchainsByAlias[alias], existingToolchain.version >= toolchain.version {
+                    continue
+                }
 
-            toolchainsByAlias[alias] = toolchain
+                state.toolchainsByAlias[alias] = toolchain
+            }
         }
     }
 
     /// Look up the toolchain with the given identifier.
     public func lookup(_ identifier: String) -> Toolchain? {
-        let lowercasedIdentifier = identifier.lowercased()
-        if ["default", "xcode"].contains(lowercasedIdentifier) {
-            return toolchainsByIdentifier[ToolchainRegistry.defaultToolchainIdentifier] ?? toolchainsByAlias[lowercasedIdentifier]
-        } else {
-            return toolchainsByIdentifier[identifier] ?? toolchainsByAlias[lowercasedIdentifier]
+        state.withLock { state in
+            let lowercasedIdentifier = identifier.lowercased()
+            if ["default", "xcode"].contains(lowercasedIdentifier) {
+                return state.toolchainsByIdentifier[ToolchainRegistry.defaultToolchainIdentifier] ?? state.toolchainsByAlias[lowercasedIdentifier]
+            } else {
+                return state.toolchainsByIdentifier[identifier] ?? state.toolchainsByAlias[lowercasedIdentifier]
+            }
         }
     }
 
@@ -545,6 +570,8 @@ public final class ToolchainRegistry: @unchecked Sendable {
     }
 
     public var toolchains: Set<Toolchain> {
-        return Set(self.toolchainsByIdentifier.values)
+        state.withLock { state in
+            return Set(state.toolchainsByIdentifier.values)
+        }
     }
 }

Sources/SWBProtocol/Message.swift

@@ -356,6 +356,20 @@ public struct MacCatalystUnavailableFrameworkNamesRequest: RequestMessage, Equat
     }
 }
 
+public struct RegisterToolchainRequest: SessionMessage, RequestMessage, Equatable, SerializableCodable {
+    public typealias ResponseMessage = StringResponse
+
+    public static let name = "REGISTER_TOOLCHAIN_REQUEST"
+
+    public let sessionHandle: String
+    public let path: Path
+
+    public init(sessionHandle: String, path: Path) {
+        self.sessionHandle = sessionHandle
+        self.path = path
+    }
+}
+
 public struct AppleSystemFrameworkNamesRequest: RequestMessage, Equatable, PendingSerializableCodable {
     public typealias ResponseMessage = StringListResponse
 
@@ -1168,6 +1182,7 @@ public struct IPCMessage: Serializable, Sendable {
         GetSpecsRequest.self,
         GetStatisticsRequest.self,
         GetToolchainsRequest.self,
+        RegisterToolchainRequest.self,
         GetBuildSettingsDescriptionRequest.self,
         ExecuteCommandLineToolRequest.self,
 

Sources/SwiftBuild/SWBBuildServiceSession.swift

@@ -629,6 +629,10 @@ public final class SWBBuildServiceSession: Sendable {
     public func setUserPreferences(enableDebugActivityLogs: Bool, enableBuildDebugging: Bool, enableBuildSystemCaching: Bool, activityTextShorteningLevel: Int, usePerConfigurationBuildLocations: Bool?, allowsExternalToolExecution: Bool) async throws {
         _ = try await service.send(request: SetSessionUserPreferencesRequest(sessionHandle: self.uid, enableDebugActivityLogs: enableDebugActivityLogs, enableBuildDebugging: enableBuildDebugging, enableBuildSystemCaching: enableBuildSystemCaching, activityTextShorteningLevel: ActivityTextShorteningLevel(rawValue: activityTextShorteningLevel) ?? .default, usePerConfigurationBuildLocations: usePerConfigurationBuildLocations, allowsExternalToolExecution: allowsExternalToolExecution))
     }
+
+    public func registerToolchain(at path: String) async throws -> String {
+        return try await service.send(request: RegisterToolchainRequest(sessionHandle: self.uid, path: Path(path))).value
+    }
 }
 
 extension SWBBuildServiceSession {

Tests/SWBCoreTests/ToolchainRegistryTests.swift

@@ -124,7 +124,7 @@ import SWBServiceCore
             ("swift.xctoolchain", ["CFBundleIdentifier": "org.swift.3020161115a", "Aliases": ["swift"]]),
             ("swift-latest.xctoolchain", ["CFBundleIdentifier": "org.swift.latest"]),
         ]) { registry, warnings, errors in
-            #expect(registry.toolchainsByIdentifier.keys.sorted(by: <) == [ToolchainRegistry.defaultToolchainIdentifier, "d", "org.swift.3020161115a"])
+            #expect(registry.toolchains.map(\.identifier).sorted(by: <) == [ToolchainRegistry.defaultToolchainIdentifier, "d", "org.swift.3020161115a"])
 
             if strict {
                 #expect(warnings.isEmpty)
@@ -175,7 +175,7 @@ import SWBServiceCore
             ("swift-older.xctoolchain", ["CFBundleIdentifier": "org.swift.3020161114a", "Version": "3.0.220161211141", "Aliases": ["swift"]]),
         ], infoPlistName: "Info.plist") { registry, _, errors in
 
-            #expect(Set(registry.toolchainsByIdentifier.keys) == Set(["org.swift.3020161114a", "org.swift.3020161115a"] + additionalToolchains))
+            #expect(Set(registry.toolchains.map(\.identifier)) == Set(["org.swift.3020161114a", "org.swift.3020161115a"] + additionalToolchains))
             #expect(errors.count == 0, "\(errors)")
             #expect(registry.lookup("org.swift.3020161115a")?.identifier == "org.swift.3020161115a")
             #expect(registry.lookup("org.swift.3020161114a")?.identifier == "org.swift.3020161114a")

Tests/SwiftBuildTests/ToolchainTests.swift

@@ -0,0 +1,47 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift open source project
+//
+// Copyright (c) 2025 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See http://swift.org/LICENSE.txt for license information
+// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import Foundation
+import Testing
+import SwiftBuild
+import SwiftBuildTestSupport
+import SWBTestSupport
+@_spi(Testing) import SWBUtil
+
+@Suite
+fileprivate struct ToolchainTests {
+    @Test
+    func lateRegistration() async throws {
+        try await withTemporaryDirectory { temporaryDirectory in
+            try await withAsyncDeferrable { deferrable in
+                let tmpDirPath = temporaryDirectory.path
+                let testSession = try await TestSWBSession(temporaryDirectory: temporaryDirectory)
+                await deferrable.addBlock {
+                    await #expect(throws: Never.self) {
+                        try await testSession.close()
+                    }
+                }
+                try localFS.createDirectory(tmpDirPath.join("Foo.xctoolchain"))
+                try await localFS.writePlist(tmpDirPath.join("Foo.xctoolchain/Info.plist"), ["Identifier" : "org.swift.foo"])
+                do {
+                    let identifier = try await testSession.session.registerToolchain(at: tmpDirPath.join("Foo.xctoolchain").str)
+                    #expect(identifier == "org.swift.foo")
+                }
+                // Late registration should be idempotent
+                do {
+                    let identifier = try await testSession.session.registerToolchain(at: tmpDirPath.join("Foo.xctoolchain").str)
+                    #expect(identifier == "org.swift.foo")
+                }
+            }
+        }
+    }
+}