fix(amazonq): add cancel support to loading developer profiles (#940)

Problem
There can be some latency while loading developer profiles, caused by either a poor internet connection, a sluggish backend, or both. Clients need a way to cancel the request for loading profiles.

Solution

Wire in support for the Cancellation Token.
This change is being done in tandem with VS Toolkit change aws/aws-toolkit-visual-studio-staging#2443

Author: awschristou

server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.test.ts

@@ -8,7 +8,7 @@ import {
 } from './qConfigurationServer'
 import { TestFeatures } from '@aws/language-server-runtimes/testing'
 import { CodeWhispererServiceToken } from '../../shared/codeWhispererService'
-import { InitializeParams, Server } from '@aws/language-server-runtimes/server-interface'
+import { CancellationTokenSource, InitializeParams, Server } from '@aws/language-server-runtimes/server-interface'
 import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager'
 import { setCredentialsForAmazonQTokenServiceManagerFactory } from '../../shared/testUtils'
 import { Q_CONFIGURATION_SECTION } from '../../shared/constants'
@@ -96,6 +96,7 @@ describe('ServerConfigurationProvider', () => {
     let codeWhispererService: StubbedInstance<CodeWhispererServiceToken>
     let testFeatures: TestFeatures
     let listAvailableProfilesHandlerSpy: sinon.SinonSpy
+    let tokenSource: CancellationTokenSource
 
     const setCredentials = setCredentialsForAmazonQTokenServiceManagerFactory(() => testFeatures)
 
@@ -119,6 +120,7 @@ describe('ServerConfigurationProvider', () => {
     }
 
     beforeEach(() => {
+        tokenSource = new CancellationTokenSource()
         codeWhispererService = stubInterface<CodeWhispererServiceToken>()
         codeWhispererService.listAvailableCustomizations.resolves({
             customizations: [],
@@ -152,14 +154,14 @@ describe('ServerConfigurationProvider', () => {
     it(`does not use listAvailableProfiles handler when developer profiles is disabled`, async () => {
         setupServerConfigurationProvider(false)
 
-        const result = await serverConfigurationProvider.listAvailableProfiles()
+        const result = await serverConfigurationProvider.listAvailableProfiles(tokenSource.token)
 
         sinon.assert.notCalled(listAvailableProfilesHandlerSpy)
         assert.deepStrictEqual(result, [])
     })
 
     it(`uses listAvailableProfiles handler when developer profiles is enabled`, async () => {
-        await serverConfigurationProvider.listAvailableProfiles()
+        await serverConfigurationProvider.listAvailableProfiles(tokenSource.token)
 
         sinon.assert.calledOnce(listAvailableProfilesHandlerSpy)
     })

server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.ts

@@ -80,9 +80,11 @@ export const QConfigurationServerToken =
                         case Q_CONFIGURATION_SECTION:
                             ;[customizations, developerProfiles] = await Promise.all([
                                 serverConfigurationProvider.listAvailableCustomizations(),
-                                serverConfigurationProvider.listAvailableProfiles(),
+                                serverConfigurationProvider.listAvailableProfiles(token),
                             ])
 
+                            throwIfCancelled(token)
+
                             return amazonQServiceManager.getEnableDeveloperProfileSupport()
                                 ? { customizations, developerProfiles }
                                 : { customizations }
@@ -91,7 +93,9 @@ export const QConfigurationServerToken =
 
                             return customizations
                         case Q_DEVELOPER_PROFILES_CONFIGURATION_SECTION:
-                            developerProfiles = await serverConfigurationProvider.listAvailableProfiles()
+                            developerProfiles = await serverConfigurationProvider.listAvailableProfiles(token)
+
+                            throwIfCancelled(token)
 
                             return developerProfiles
                         default:
@@ -115,6 +119,12 @@ export const QConfigurationServerToken =
         return () => {}
     }
 
+function throwIfCancelled(token: CancellationToken) {
+    if (token.isCancellationRequested) {
+        throw new ResponseError(LSPErrorCodes.RequestCancelled, 'Request cancelled')
+    }
+}
+
 const ON_GET_CONFIGURATION_FROM_SERVER_ERROR_PREFIX = 'Failed to fetch: '
 
 export class ServerConfigurationProvider {
@@ -130,7 +140,7 @@ export class ServerConfigurationProvider {
         )
     }
 
-    async listAvailableProfiles(): Promise<AmazonQDeveloperProfile[]> {
+    async listAvailableProfiles(token: CancellationToken): Promise<AmazonQDeveloperProfile[]> {
         if (!this.serviceManager.getEnableDeveloperProfileSupport()) {
             this.logging.debug('Q developer profiles disabled - returning empty list')
             return []
@@ -140,6 +150,7 @@ export class ServerConfigurationProvider {
             const profiles = await this.listAllAvailableProfilesHandler({
                 connectionType: this.credentialsProvider.getConnectionType(),
                 logging: this.logging,
+                token: token,
             })
 
             return profiles

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts

@@ -305,6 +305,7 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<CodeWh
         const profiles = await getListAllAvailableProfilesHandler(this.serviceFactory)({
             connectionType: 'identityCenter',
             logging: this.logging,
+            token: token,
         })
 
         this.handleTokenCancellationRequest(token)

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/qDeveloperProfiles.test.ts

@@ -2,7 +2,11 @@ import * as assert from 'assert'
 import sinon, { StubbedInstance, stubInterface } from 'ts-sinon'
 import { CodeWhispererServiceToken } from '../codeWhispererService'
 import { SsoConnectionType } from '../utils'
-import { AWSInitializationOptions, Logging } from '@aws/language-server-runtimes/server-interface'
+import {
+    AWSInitializationOptions,
+    CancellationTokenSource,
+    Logging,
+} from '@aws/language-server-runtimes/server-interface'
 import {
     AmazonQDeveloperProfile,
     getListAllAvailableProfilesHandler,
@@ -35,6 +39,7 @@ describe('ListAllAvailableProfiles Handler', () => {
 
     let codeWhispererService: StubbedInstance<CodeWhispererServiceToken>
     let handler: ListAllAvailableProfilesHandler
+    let tokenSource: CancellationTokenSource
 
     const listAvailableProfilesResponse = {
         profiles: [
@@ -52,13 +57,15 @@ describe('ListAllAvailableProfiles Handler', () => {
         codeWhispererService.listAvailableProfiles.resolves(listAvailableProfilesResponse)
 
         handler = getListAllAvailableProfilesHandler(() => codeWhispererService)
+        tokenSource = new CancellationTokenSource()
     })
 
     it('should aggregrate profiles retrieved from different regions', async () => {
         const profiles = await handler({
             connectionType: 'identityCenter',
             logging,
             endpoints: SOME_AWS_Q_ENDPOINTS,
+            token: tokenSource.token,
         })
 
         assert.strictEqual(
@@ -73,6 +80,7 @@ describe('ListAllAvailableProfiles Handler', () => {
             const profiles = await handler({
                 connectionType,
                 logging,
+                token: tokenSource.token,
             })
 
             assert.deepStrictEqual(profiles, [])
@@ -103,6 +111,7 @@ describe('ListAllAvailableProfiles Handler', () => {
                 connectionType: 'identityCenter',
                 logging,
                 endpoints: SOME_AWS_Q_ENDPOINT,
+                token: tokenSource.token,
             })
 
             sinon.assert.calledThrice(codeWhispererService.listAvailableProfiles)
@@ -120,6 +129,7 @@ describe('ListAllAvailableProfiles Handler', () => {
                 connectionType: 'identityCenter',
                 logging,
                 endpoints: SOME_AWS_Q_ENDPOINT,
+                token: tokenSource.token,
             })
 
             assert.strictEqual(codeWhispererService.listAvailableProfiles.callCount, MAX_EXPECTED_PAGES)

server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/qDeveloperProfiles.ts

@@ -1,5 +1,6 @@
 import {
     AWSInitializationOptions,
+    CancellationToken,
     Logging,
     LSPErrorCodes,
     ResponseError,
@@ -22,6 +23,7 @@ export interface ListAllAvailableProfilesHandlerParams {
     connectionType: SsoConnectionType
     logging: Logging
     endpoints?: Map<string, string> // override option for flexibility, we default to all (AWS_Q_ENDPOINTS)
+    token: CancellationToken
 }
 
 export type ListAllAvailableProfilesHandler = (
@@ -33,7 +35,7 @@ const MAX_Q_DEVELOPER_PROFILES_PER_PAGE = 10
 
 export const getListAllAvailableProfilesHandler =
     (service: (region: string, endpoint: string) => CodeWhispererServiceToken): ListAllAvailableProfilesHandler =>
-    async ({ connectionType, logging, endpoints }) => {
+    async ({ connectionType, logging, endpoints, token }) => {
         if (!connectionType || connectionType !== 'identityCenter') {
             logging.debug('Connection type is not set or not identityCenter - returning empty response.')
             return []
@@ -42,13 +44,21 @@ export const getListAllAvailableProfilesHandler =
         let allProfiles: AmazonQDeveloperProfile[] = []
         const qEndpoints = endpoints ?? AWS_Q_ENDPOINTS
 
+        if (token.isCancellationRequested) {
+            return []
+        }
+
         const result = await Promise.allSettled(
             Array.from(qEndpoints.entries(), ([region, endpoint]) => {
                 const codeWhispererService = service(region, endpoint)
-                return fetchProfilesFromRegion(codeWhispererService, region, logging)
+                return fetchProfilesFromRegion(codeWhispererService, region, logging, token)
             })
         )
 
+        if (token.isCancellationRequested) {
+            return []
+        }
+
         const fulfilledResults = result.filter(settledResult => settledResult.status === 'fulfilled')
 
         if (fulfilledResults.length === 0) {
@@ -63,7 +73,8 @@ export const getListAllAvailableProfilesHandler =
 async function fetchProfilesFromRegion(
     service: CodeWhispererServiceToken,
     region: string,
-    logging: Logging
+    logging: Logging,
+    token: CancellationToken
 ): Promise<AmazonQDeveloperProfile[]> {
     let allRegionalProfiles: AmazonQDeveloperProfile[] = []
     let nextToken: string | undefined = undefined
@@ -73,6 +84,10 @@ async function fetchProfilesFromRegion(
         do {
             logging.debug(`Fetching profiles from region: ${region} (iteration: ${numberOfPages})`)
 
+            if (token.isCancellationRequested) {
+                return allRegionalProfiles
+            }
+
             const response = await service.listAvailableProfiles({
                 maxResults: MAX_Q_DEVELOPER_PROFILES_PER_PAGE,
                 nextToken: nextToken,