Ensure folders are added to sourcekit-lsp in multi root workspaces

plemarquand · plemarquand · commit f95c1e9de85a · 2025-06-30T16:04:33.000-04:00
`addFolder` in the LSP client was checking to see if the folder being added was a root folder. If the user is working in a multi root workspace document symbols would not be provided if they were for any folder except the first one added. This also prevented the test explorer from properly initializing. Check to see if the folder being added is in the list of VS Code workspace folders and associate the folder with the onDocumentSymbols request from the LSP. Issue: #1669
diff --git a/src/FolderContext.ts b/src/FolderContext.ts
@@ -177,9 +177,9 @@ export class FolderContext implements vscode.Disposable {
     }
 
     /** Refresh the tests in the test explorer for this folder */
-    refreshTestExplorer() {
+    async refreshTestExplorer() {
         if (this.testExplorer?.controller.resolveHandler) {
-            void this.testExplorer.controller.resolveHandler(undefined);
+            return this.testExplorer.controller.resolveHandler(undefined);
         }
     }
 
@@ -211,6 +211,25 @@ export class FolderContext implements vscode.Disposable {
         });
         return target;
     }
+
+    /**
+     * Called whenever we have new document symbols
+     */
+    onDocumentSymbols(
+        document: vscode.TextDocument,
+        symbols: vscode.DocumentSymbol[] | null | undefined
+    ) {
+        const uri = document?.uri;
+        if (
+            this.testExplorer &&
+            symbols &&
+            uri &&
+            uri.scheme === "file" &&
+            isPathInsidePath(uri.fsPath, this.folder.fsPath)
+        ) {
+            void this.testExplorer.getDocumentTests(this, uri, symbols);
+        }
+    }
 }
 
 export interface EditedPackage {
diff --git a/src/TestExplorer/TestExplorer.ts b/src/TestExplorer/TestExplorer.ts
@@ -15,7 +15,6 @@
 import * as vscode from "vscode";
 import { FolderContext } from "../FolderContext";
 import { getErrorDescription } from "../utilities/utilities";
-import { isPathInsidePath } from "../utilities/filesystem";
 import { FolderOperation, WorkspaceContext } from "../WorkspaceContext";
 import { TestRunProxy, TestRunner } from "./TestRunner";
 import { LSPTestDiscovery } from "./LSPTestDiscovery";
@@ -136,45 +135,11 @@ export class TestExplorer {
         const disposable = workspaceContext.onDidChangeFolders(({ folder, operation }) => {
             switch (operation) {
                 case FolderOperation.add:
-                    if (folder) {
-                        void folder.swiftPackage.getTargets(TargetType.test).then(targets => {
-                            if (targets.length === 0) {
-                                return;
-                            }
-
-                            folder.addTestExplorer();
-                            // discover tests in workspace but only if disableAutoResolve is not on.
-                            // discover tests will kick off a resolve if required
-                            if (!configuration.folder(folder.workspaceFolder).disableAutoResolve) {
-                                void folder.testExplorer?.discoverTestsInWorkspace(
-                                    tokenSource.token
-                                );
-                            }
-                        });
-                    }
-                    break;
                 case FolderOperation.packageUpdated:
                     if (folder) {
-                        void folder.swiftPackage.getTargets(TargetType.test).then(targets => {
-                            const hasTestTargets = targets.length > 0;
-                            if (hasTestTargets && !folder.hasTestExplorer()) {
-                                folder.addTestExplorer();
-                                // discover tests in workspace but only if disableAutoResolve is not on.
-                                // discover tests will kick off a resolve if required
-                                if (
-                                    !configuration.folder(folder.workspaceFolder).disableAutoResolve
-                                ) {
-                                    void folder.testExplorer?.discoverTestsInWorkspace(
-                                        tokenSource.token
-                                    );
-                                }
-                            } else if (!hasTestTargets && folder.hasTestExplorer()) {
-                                folder.removeTestExplorer();
-                            } else if (folder.hasTestExplorer()) {
-                                folder.refreshTestExplorer();
-                            }
-                        });
+                        void this.setupTestExplorerForFolder(folder, tokenSource.token);
                     }
+                    break;
             }
         });
         return {
@@ -185,6 +150,28 @@ export class TestExplorer {
         };
     }
 
+    /**
+     * Configures a test explorer for the given folder.
+     * If the folder has test targets, and there is no existing test explorer,
+     * it will create a test explorer and discover tests.
+     * If the folder has no test targets, it will remove any existing test explorer.
+     * If the folder has test targets and an existing test explorer, it will refresh the tests.
+     */
+    private static async setupTestExplorerForFolder(
+        folder: FolderContext,
+        token: vscode.CancellationToken
+    ) {
+        const targets = await folder.swiftPackage.getTargets(TargetType.test);
+        const hasTestTargets = targets.length > 0;
+        if (hasTestTargets && !folder.hasTestExplorer()) {
+            await folder.addTestExplorer().discoverTestsInWorkspace(token);
+        } else if (hasTestTargets && folder.hasTestExplorer()) {
+            await folder.refreshTestExplorer();
+        } else if (!hasTestTargets && folder.hasTestExplorer()) {
+            folder.removeTestExplorer();
+        }
+    }
+
     /**
      * Sets the `swift.tests` context variable which is used by commands
      * to determine if the test item belongs to the Swift extension.
@@ -196,45 +183,29 @@ export class TestExplorer {
         });
     }
 
-    /**
-     * Called whenever we have new document symbols
-     */
-    static onDocumentSymbols(
+    async getDocumentTests(
         folder: FolderContext,
-        document: vscode.TextDocument,
-        symbols: vscode.DocumentSymbol[] | null | undefined
-    ) {
-        const uri = document?.uri;
-        const testExplorer = folder?.testExplorer;
-        if (testExplorer && symbols && uri && uri.scheme === "file") {
-            if (isPathInsidePath(uri.fsPath, folder.folder.fsPath)) {
-                void folder.swiftPackage.getTarget(uri.fsPath).then(target => {
-                    if (target && target.type === "test") {
-                        testExplorer.lspTestDiscovery
-                            .getDocumentTests(folder.swiftPackage, uri)
-                            .then(tests => {
-                                TestDiscovery.updateTestsForTarget(
-                                    testExplorer.controller,
-                                    { id: target.c99name, label: target.name },
-                                    tests,
-                                    uri
-                                );
-                                testExplorer.onTestItemsDidChangeEmitter.fire(
-                                    testExplorer.controller
-                                );
-                            })
-                            // Fallback to parsing document symbols for XCTests only
-                            .catch(() => {
-                                const tests = parseTestsFromDocumentSymbols(
-                                    target.name,
-                                    symbols,
-                                    uri
-                                );
-                                testExplorer.updateTests(testExplorer.controller, tests, uri);
-                            });
-                    }
-                });
-            }
+        uri: vscode.Uri,
+        symbols: vscode.DocumentSymbol[]
+    ): Promise<void> {
+        const target = await folder.swiftPackage.getTarget(uri.fsPath);
+        if (!target || target.type !== "test") {
+            return;
+        }
+
+        try {
+            const tests = await this.lspTestDiscovery.getDocumentTests(folder.swiftPackage, uri);
+            TestDiscovery.updateTestsForTarget(
+                this.controller,
+                { id: target.c99name, label: target.name },
+                tests,
+                uri
+            );
+            this.onTestItemsDidChangeEmitter.fire(this.controller);
+        } catch {
+            // Fallback to parsing document symbols for XCTests only
+            const tests = parseTestsFromDocumentSymbols(target.name, symbols, uri);
+            this.updateTests(this.controller, tests, uri);
         }
     }
 
diff --git a/src/WorkspaceContext.ts b/src/WorkspaceContext.ts
@@ -37,7 +37,6 @@ import { isValidWorkspaceFolder, searchForPackages } from "./utilities/workspace
 import { SwiftPluginTaskProvider } from "./tasks/SwiftPluginTaskProvider";
 import { SwiftTaskProvider } from "./tasks/SwiftTaskProvider";
 import { LLDBDebugConfigurationProvider } from "./debugger/debugAdapterFactory";
-import { TestExplorer } from "./TestExplorer/TestExplorer";
 
 /**
  * Context for whole workspace. Holds array of contexts for each workspace folder
@@ -82,7 +81,7 @@ export class WorkspaceContext implements vscode.Disposable {
         this.buildStatus = new SwiftBuildStatus(this.statusItem);
         this.languageClientManager = new LanguageClientToolchainCoordinator(this, {
             onDocumentSymbols: (folder, document, symbols) => {
-                TestExplorer.onDocumentSymbols(folder, document, symbols);
+                folder.onDocumentSymbols(document, symbols);
             },
         });
         this.tasks = new TaskManager(this);
diff --git a/src/sourcekit-lsp/LanguageClientManager.ts b/src/sourcekit-lsp/LanguageClientManager.ts
@@ -97,7 +97,8 @@ export class LanguageClientManager implements vscode.Disposable {
     private singleServerSupport: boolean;
     // used by single server support to keep a record of the project folders
     // that are not at the root of their workspace
-    public subFolderWorkspaces: vscode.Uri[];
+    public subFolderWorkspaces: FolderContext[] = [];
+    private addedFolders: FolderContext[] = [];
     private namedOutputChannels: Map<string, LSPOutputChannel> = new Map();
     private swiftVersion: Version;
     private activeDocumentManager = new LSPActiveDocumentManager();
@@ -122,7 +123,6 @@ export class LanguageClientManager implements vscode.Disposable {
         this.swiftVersion = folderContext.swiftVersion;
         this.singleServerSupport = this.swiftVersion.isGreaterThanOrEqual(new Version(5, 7, 0));
         this.subscriptions = [];
-        this.subFolderWorkspaces = [];
 
         // on change config restart server
         const onChangeConfig = vscode.workspace.onDidChangeConfiguration(event => {
@@ -244,9 +244,9 @@ export class LanguageClientManager implements vscode.Disposable {
     async addFolder(folderContext: FolderContext) {
         if (!folderContext.isRootFolder) {
             await this.useLanguageClient(async client => {
-                const uri = folderContext.folder;
-                this.subFolderWorkspaces.push(folderContext.folder);
+                this.subFolderWorkspaces.push(folderContext);
 
+                const uri = folderContext.folder;
                 const workspaceFolder = {
                     uri: client.code2ProtocolConverter.asUri(uri),
                     name: FolderContext.uriName(uri),
@@ -256,13 +256,16 @@ export class LanguageClientManager implements vscode.Disposable {
                 });
             });
         }
+        this.addedFolders.push(folderContext);
     }
 
     async removeFolder(folderContext: FolderContext) {
         if (!folderContext.isRootFolder) {
             await this.useLanguageClient(async client => {
                 const uri = folderContext.folder;
-                this.subFolderWorkspaces = this.subFolderWorkspaces.filter(item => item !== uri);
+                this.subFolderWorkspaces = this.subFolderWorkspaces.filter(
+                    item => item.folder !== uri
+                );
 
                 const workspaceFolder = {
                     uri: client.code2ProtocolConverter.asUri(uri),
@@ -273,13 +276,14 @@ export class LanguageClientManager implements vscode.Disposable {
                 });
             });
         }
+        this.addedFolders = this.addedFolders.filter(item => item.folder !== folderContext.folder);
     }
 
     private async addSubFolderWorkspaces(client: LanguageClient) {
-        for (const uri of this.subFolderWorkspaces) {
+        for (const folderContext of this.subFolderWorkspaces) {
             const workspaceFolder = {
-                uri: client.code2ProtocolConverter.asUri(uri),
-                name: FolderContext.uriName(uri),
+                uri: client.code2ProtocolConverter.asUri(folderContext.folder),
+                name: FolderContext.uriName(folderContext.folder),
             };
             await client.sendNotification(DidChangeWorkspaceFoldersNotification.type, {
                 event: { added: [workspaceFolder], removed: [] },
@@ -440,7 +444,17 @@ export class LanguageClientManager implements vscode.Disposable {
             this.activeDocumentManager,
             errorHandler,
             (document, symbols) => {
-                this.options.onDocumentSymbols?.(this.folderContext, document, symbols);
+                const documentFolderContext = [this.folderContext, ...this.addedFolders].find(
+                    folderContext => document.uri.fsPath.startsWith(folderContext.folder.fsPath)
+                );
+                if (!documentFolderContext) {
+                    this.languageClientOutputChannel?.log(
+                        "Unable to find folder for document: " + document.uri.fsPath,
+                        "WARN"
+                    );
+                    return;
+                }
+                this.options.onDocumentSymbols?.(documentFolderContext, document, symbols);
             }
         );
 
diff --git a/test/integration-tests/ExtensionActivation.test.ts b/test/integration-tests/ExtensionActivation.test.ts
@@ -116,7 +116,9 @@ suite("Extension Activation/Deactivation Tests", () => {
             assert(folder);
 
             const languageClient = workspaceContext.languageClientManager.get(folder);
-            const lspWorkspaces = languageClient.subFolderWorkspaces.map(({ fsPath }) => fsPath);
+            const lspWorkspaces = languageClient.subFolderWorkspaces.map(
+                ({ folder }) => folder.fsPath
+            );
             assertContains(lspWorkspaces, testAssetUri("cmake").fsPath);
         });
 
@@ -125,7 +127,9 @@ suite("Extension Activation/Deactivation Tests", () => {
             assert(folder);
 
             const languageClient = workspaceContext.languageClientManager.get(folder);
-            const lspWorkspaces = languageClient.subFolderWorkspaces.map(({ fsPath }) => fsPath);
+            const lspWorkspaces = languageClient.subFolderWorkspaces.map(
+                ({ folder }) => folder.fsPath
+            );
             assertContains(lspWorkspaces, testAssetUri("cmake-compile-flags").fsPath);
         });
     });
diff --git a/test/integration-tests/commands/runTestMultipleTimes.test.ts b/test/integration-tests/commands/runTestMultipleTimes.test.ts
@@ -25,7 +25,7 @@ suite("Test Multiple Times Command Test Suite", () => {
 
     activateExtensionForSuite({
         async setup(ctx) {
-            folderContext = await folderInRootWorkspace("diagnostics", ctx);
+            folderContext = await folderInRootWorkspace("defaultPackage", ctx);
             folderContext.addTestExplorer();
 
             const item = folderContext.testExplorer?.controller.createTestItem(
diff --git a/test/integration-tests/testexplorer/TestExplorerIntegration.test.ts b/test/integration-tests/testexplorer/TestExplorerIntegration.test.ts
