Await Workspace Trust transition completion in setUrisTrust()

src/vs/workbench/services/workspaces/common/workspaceTrust.ts

@@ -625,7 +625,7 @@ export class WorkspaceTrustManagementService extends Disposable implements IWork
 	}
 
 	async setUrisTrust(uris: URI[], trusted: boolean): Promise<void> {
-		this.doSetUrisTrust(await Promise.all(uris.map(uri => this.getCanonicalUri(uri))), trusted);
+		await this.doSetUrisTrust(await Promise.all(uris.map(uri => this.getCanonicalUri(uri))), trusted);
 	}
 
 	getTrustedUris(): URI[] {

src/vs/workbench/services/workspaces/test/common/workspaceTrust.test.ts

@@ -13,7 +13,7 @@ import { TestInstantiationService } from '../../../../../platform/instantiation/
 import { NullLogService } from '../../../../../platform/log/common/log.js';
 import { IRemoteAuthorityResolverService } from '../../../../../platform/remote/common/remoteAuthorityResolver.js';
 import { IStorageService, StorageScope, StorageTarget } from '../../../../../platform/storage/common/storage.js';
-import { IWorkspaceContextService } from '../../../../../platform/workspace/common/workspace.js';
+import { IWorkspaceContextService, toWorkspaceFolder } from '../../../../../platform/workspace/common/workspace.js';
 import { IWorkspaceTrustEnablementService, IWorkspaceTrustInfo } from '../../../../../platform/workspace/common/workspaceTrust.js';
 import { Workspace } from '../../../../../platform/workspace/test/common/testWorkspace.js';
 import { Memento } from '../../../../common/memento.js';
@@ -159,6 +159,96 @@ suite('Workspace Trust', () => {
 			assert.strictEqual(true, (await testObject.getUriTrustInfo(sameFolderDifferentMeta)).trusted);
 		});
 
+		test('setWorkspaceTrust waits for trust transition participants before resolving', async () => {
+			await configurationService.setUserConfiguration('security', getUserSettings(true, true));
+			workspaceService.setWorkspace(new Workspace('folder-workspace', [toWorkspaceFolder(URI.parse('file:///Folder'))]));
+			const testObject = await initializeTestObject();
+
+			let releaseParticipant!: () => void;
+			const participantCanComplete = new Promise<void>(resolve => releaseParticipant = resolve);
+
+			let participantStartedResolve!: () => void;
+			const participantStarted = new Promise<void>(resolve => participantStartedResolve = resolve);
+
+			let participantStartedFlag = false;
+			let participantCompleted = false;
+			let trustChangeEventFired = false;
+
+			const participantCompletedPromise = new Promise<void>(resolve => {
+				store.add(testObject.addWorkspaceTrustTransitionParticipant({
+					async participate(trusted: boolean): Promise<void> {
+						if (trusted) {
+							participantStartedFlag = true;
+							participantStartedResolve();
+							await participantCanComplete;
+							participantCompleted = true;
+							resolve();
+						}
+					}
+				}));
+			});
+
+			store.add(testObject.onDidChangeTrust(trusted => {
+				if (trusted) {
+					trustChangeEventFired = true;
+				}
+			}));
+
+			await testObject.setWorkspaceTrust(false);
+			assert.deepStrictEqual({
+				trusted: testObject.isWorkspaceTrusted(),
+				participantStarted: participantStartedFlag,
+				participantCompleted,
+				trustChangeEventFired
+			}, {
+				trusted: false,
+				participantStarted: false,
+				participantCompleted: false,
+				trustChangeEventFired: false
+			});
+
+			const setWorkspaceTrustPromise = testObject.setWorkspaceTrust(true);
+			let setWorkspaceTrustResolved = false;
+			setWorkspaceTrustPromise.then(() => setWorkspaceTrustResolved = true);
+
+			try {
+				await participantStarted;
+				await Promise.resolve();
+
+				assert.deepStrictEqual({
+					setWorkspaceTrustResolved,
+					trusted: testObject.isWorkspaceTrusted(),
+					participantStarted: participantStartedFlag,
+					participantCompleted,
+					trustChangeEventFired
+				}, {
+					setWorkspaceTrustResolved: false,
+					trusted: true,
+					participantStarted: true,
+					participantCompleted: false,
+					trustChangeEventFired: false
+				});
+			} finally {
+				releaseParticipant();
+				await participantCompletedPromise;
+			}
+
+			await setWorkspaceTrustPromise;
+			await Promise.resolve();
+
+			assert.deepStrictEqual({
+				setWorkspaceTrustResolved,
+				trusted: testObject.isWorkspaceTrusted(),
+				participantCompleted,
+				trustChangeEventFired
+			}, {
+				setWorkspaceTrustResolved: true,
+				trusted: true,
+				participantCompleted: true,
+				trustChangeEventFired: true
+			});
+		});
+
 		async function initializeTestObject(): Promise<WorkspaceTrustManagementService> {
 			const workspaceTrustManagementService = store.add(instantiationService.createInstance(WorkspaceTrustManagementService));
 			await workspaceTrustManagementService.workspaceTrustInitialized;