Allow for querying multiple kinds of related information

src/vs/workbench/services/aiRelatedInformation/common/aiRelatedInformationService.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { CancellationToken } from 'vs/base/common/cancellation';
-import { CancelablePromise, createCancelablePromise, raceCancellablePromises, timeout } from 'vs/base/common/async';
+import { CancelablePromise, createCancelablePromise, raceTimeout } from 'vs/base/common/async';
 import { IDisposable } from 'vs/base/common/lifecycle';
 import { InstantiationType, registerSingleton } from 'vs/platform/instantiation/common/extensions';
 import { StopWatch } from 'vs/base/common/stopwatch';
@@ -64,42 +64,35 @@ export class AiRelatedInformationService implements IAiRelatedInformationService
 
 		const stopwatch = StopWatch.create();
 
-		const cancellablePromises: Array<CancelablePromise<RelatedInformationResult[]>> = [];
-
-		const timer = timeout(AiRelatedInformationService.DEFAULT_TIMEOUT);
-		const disposable = token.onCancellationRequested(() => {
-			disposable.dispose();
-			timer.cancel();
-		});
-
-		for (const provider of providers) {
-			cancellablePromises.push(createCancelablePromise(async t => {
+		const cancellablePromises: Array<CancelablePromise<RelatedInformationResult[]>> = providers.map((provider) => {
+			return createCancelablePromise(async t => {
 				try {
 					const result = await provider.provideAiRelatedInformation(query, t);
 					// double filter just in case
 					return result.filter(r => types.includes(r.type));
 				} catch (e) {
 					// logged in extension host
 				}
-				// Wait for the timer to finish to allow for another provider to resolve.
-				// Alternatively, if something resolved, or we've timed out, this will throw
-				// as expected.
-				await timer;
-				throw new Error('Related information provider timed out');
-			}));
-		}
-
-		cancellablePromises.push(createCancelablePromise(async (t) => {
-			const disposable = t.onCancellationRequested(() => {
-				timer.cancel();
-				disposable.dispose();
+				return [];
 			});
-			await timer;
-			throw new Error('Related information provider timed out');
-		}));
+		});
 
 		try {
-			const result = await raceCancellablePromises(cancellablePromises);
+			const results = await raceTimeout(
+				Promise.allSettled(cancellablePromises),
+				AiRelatedInformationService.DEFAULT_TIMEOUT,
+				() => {
+					cancellablePromises.forEach(p => p.cancel());
+					throw new Error('Related information provider timed out');
+				}
+			);
+			if (!results) {
+				return [];
+			}
+			const result = results
+				.filter(r => r.status === 'fulfilled')
+				.map(r => (r as PromiseFulfilledResult<RelatedInformationResult[]>).value)
+				.flat();
 			return result;
 		} finally {
 			stopwatch.stop();

src/vs/workbench/services/aiRelatedInformation/test/common/aiRelatedInformationService.test.ts

@@ -6,19 +6,21 @@
 import * as assert from 'assert';
 import { AiRelatedInformationService } from 'vs/workbench/services/aiRelatedInformation/common/aiRelatedInformationService';
 import { NullLogService } from 'vs/platform/log/common/log';
-import { CommandInformationResult, IAiRelatedInformationProvider, RelatedInformationType } from 'vs/workbench/services/aiRelatedInformation/common/aiRelatedInformation';
+import { CommandInformationResult, IAiRelatedInformationProvider, RelatedInformationType, SettingInformationResult } from 'vs/workbench/services/aiRelatedInformation/common/aiRelatedInformation';
 import { CancellationToken } from 'vs/base/common/cancellation';
+import { ensureNoDisposablesAreLeakedInTestSuite } from 'vs/base/test/common/utils';
 
 suite('AiRelatedInformationService', () => {
+	const store = ensureNoDisposablesAreLeakedInTestSuite();
 	let service: AiRelatedInformationService;
 
 	setup(() => {
-		service = new AiRelatedInformationService(new NullLogService());
+		service = new AiRelatedInformationService(store.add(new NullLogService()));
 	});
 
 	test('should check if providers are registered', () => {
 		assert.equal(service.isEnabled(), false);
-		service.registerAiRelatedInformationProvider(RelatedInformationType.CommandInformation, { provideAiRelatedInformation: () => Promise.resolve([]) });
+		store.add(service.registerAiRelatedInformationProvider(RelatedInformationType.CommandInformation, { provideAiRelatedInformation: () => Promise.resolve([]) }));
 		assert.equal(service.isEnabled(), true);
 	});
 
@@ -40,4 +42,28 @@ suite('AiRelatedInformationService', () => {
 		assert.strictEqual(result.length, 1);
 		assert.strictEqual((result[0] as CommandInformationResult).command, command);
 	});
+
+	test('should get different types of related information', async () => {
+		const command = 'command';
+		const commandProvider: IAiRelatedInformationProvider = {
+			provideAiRelatedInformation: () => Promise.resolve([{ type: RelatedInformationType.CommandInformation, command, weight: 1 }])
+		};
+		service.registerAiRelatedInformationProvider(RelatedInformationType.CommandInformation, commandProvider);
+		const setting = 'setting';
+		const settingProvider: IAiRelatedInformationProvider = {
+			provideAiRelatedInformation: () => Promise.resolve([{ type: RelatedInformationType.SettingInformation, setting, weight: 1 }])
+		};
+		service.registerAiRelatedInformationProvider(RelatedInformationType.SettingInformation, settingProvider);
+		const result = await service.getRelatedInformation(
+			'query',
+			[
+				RelatedInformationType.CommandInformation,
+				RelatedInformationType.SettingInformation
+			],
+			CancellationToken.None
+		);
+		assert.strictEqual(result.length, 2);
+		assert.strictEqual((result[0] as CommandInformationResult).command, command);
+		assert.strictEqual((result[1] as SettingInformationResult).setting, setting);
+	});
 });