From da0c9a9ad97055176d9a5178a97c94de11d35231 Mon Sep 17 00:00:00 2001
From: Bri Peticca <briannap_22@apple.com>
Date: Wed, 27 May 2026 12:57:16 -0400
Subject: [PATCH] Fix spm API usage for appropriate trait configuration

The `WorkspaceConfiguration.default` was being used to populate
the `Workspace`; there was an extra step here needed to assure that
we are propagating the trait configuration to the workspace.

The added test assures that non-default traits that are enabled are
indeed processed as enabled.
---
 .../SwiftPMBuildServer.swift                  | 19 +++++-----
 .../SwiftPMBuildServerTests.swift             | 36 +++++++++++++++++++
 2 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/Sources/BuildServerIntegration/SwiftPMBuildServer.swift b/Sources/BuildServerIntegration/SwiftPMBuildServer.swift
index 9a814ab8..d81e3d1b 100644
--- a/Sources/BuildServerIntegration/SwiftPMBuildServer.swift
+++ b/Sources/BuildServerIntegration/SwiftPMBuildServer.swift
@@ -331,8 +331,17 @@ package actor SwiftPMBuildServer: BuiltInBuildServer {
       location.scratchDirectory = absProjectRoot.appending(components: ".build", "index-build")
     }
 
+    let enabledTraits: Set<String>? =
+      if let traits = options.swiftPMOrDefault.traits {
+        Set(traits)
+      } else {
+        nil
+      }
+    self.traitConfiguration = TraitConfiguration(enabledTraits: enabledTraits)
+
     var configuration = WorkspaceConfiguration.default
     configuration.skipDependenciesUpdates = !options.backgroundIndexingOrDefault
+    configuration.traitConfiguration = self.traitConfiguration
 
     self.swiftPMWorkspace = try Workspace(
       fileSystem: localFileSystem,
@@ -394,14 +403,6 @@ package actor SwiftPMBuildServer: BuiltInBuildServer {
       disableSandbox: options.swiftPMOrDefault.disableSandbox ?? false
     )
 
-    let enabledTraits: Set<String>? =
-      if let traits = options.swiftPMOrDefault.traits {
-        Set(traits)
-      } else {
-        nil
-      }
-    self.traitConfiguration = TraitConfiguration(enabledTraits: enabledTraits)
-
     packageLoadingQueue.async {
       await orLog("Initial package loading") {
         // Schedule an initial generation of the build graph. Once the build graph is loaded, the build server will send
@@ -499,7 +500,7 @@ package actor SwiftPMBuildServer: BuiltInBuildServer {
     }
 
     let modulesGraph = try await self.swiftPMWorkspace.loadPackageGraph(
-      rootInput: PackageGraphRootInput(packages: [AbsolutePath(validating: projectRoot.filePath)]),
+      rootInput: PackageGraphRootInput(packages: [AbsolutePath(validating: projectRoot.filePath)], traitConfiguration: self.traitConfiguration),
       forceResolvedVersions: options.swiftPMOrDefault.forceResolvedVersions ?? !isForIndexBuild,
       observabilityScope: observabilitySystem.topScope.makeChildScope(description: "Load package graph")
     )
diff --git a/Tests/BuildServerIntegrationTests/SwiftPMBuildServerTests.swift b/Tests/BuildServerIntegrationTests/SwiftPMBuildServerTests.swift
index 721d4c1a..6e9c7e47 100644
--- a/Tests/BuildServerIntegrationTests/SwiftPMBuildServerTests.swift
+++ b/Tests/BuildServerIntegrationTests/SwiftPMBuildServerTests.swift
@@ -1643,6 +1643,42 @@ struct SwiftPMBuildServerTests {
       }
     }
   }
+
+  @Test
+  func testPackagePlugin() async throws {
+    let testProject = try await SwiftPMTestProject(
+      files: [
+        "Test.swift": """
+        #if NonDefaultTrait
+        #warning("Trait enabled")
+        #endif
+        """
+      ],
+      manifest: """
+        // swift-tools-version: 6.2
+
+        import PackageDescription
+
+        let package = Package(
+          name: "MyLibrary",
+          traits: [
+            .default(enabledTraits: []),
+            "NonDefaultTrait",
+          ],
+          targets: [.target(name: "MyLibrary")]
+        )
+        """,
+      options: SourceKitLSPOptions(swiftPM: .init(traits: ["NonDefaultTrait"])),
+      enableBackgroundIndexing: true
+    )
+
+    let (uri, _) = try testProject.openDocument("Test.swift")
+    let diagnostics = try await testProject.testClient.send(
+      DocumentDiagnosticsRequest(textDocument: TextDocumentIdentifier(uri))
+    )
+
+    #expect(diagnostics.fullReport?.items.map(\.message) == ["Trait enabled"])
+  }
 }
 
 private func expectArgumentsDoNotContain(