diff --git a/src/vs/workbench/services/aiRelatedInformation/common/aiRelatedInformationService.ts b/src/vs/workbench/services/aiRelatedInformation/common/aiRelatedInformationService.ts index 7168c3b4eed..cd4e247b59e 100644 --- a/src/vs/workbench/services/aiRelatedInformation/common/aiRelatedInformationService.ts +++ b/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,16 +64,8 @@ export class AiRelatedInformationService implements IAiRelatedInformationService const stopwatch = StopWatch.create(); - const cancellablePromises: Array> = []; - - 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> = providers.map((provider) => { + return createCancelablePromise(async t => { try { const result = await provider.provideAiRelatedInformation(query, t); // double filter just in case @@ -81,25 +73,26 @@ export class AiRelatedInformationService implements IAiRelatedInformationService } 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).value) + .flat(); return result; } finally { stopwatch.stop(); diff --git a/src/vs/workbench/services/aiRelatedInformation/test/common/aiRelatedInformationService.test.ts b/src/vs/workbench/services/aiRelatedInformation/test/common/aiRelatedInformationService.test.ts index 6d521ca8777..62f25baed57 100644 --- a/src/vs/workbench/services/aiRelatedInformation/test/common/aiRelatedInformationService.test.ts +++ b/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); + }); });