Skip to content

Commit 92a0710

Browse files
Will-ShaoHuaWill-ShaoHua
authored
refactor(amazonq): optimization: remove validation call for selected profile (#1329)
* refactor(amazonq): optimization: only call listAvailableProfile with region associated with the arn * revert(amazonq): console.log * refactor(amazonq): stop validating profile availability on configurationChanged * fix(amazonq): fix typo * fix(amazonq): test * test(amazonq): amazonQTokenServiceManager.test.ts * refactor(amazonq): revert unneeded changes * refactor(amazonq): nit * refactor(amazonq): more nit * refactor(amazonq): fix test * refactor(amazonq): fix test * refactor(amazonq): nit * refactor(amazonq): revert unneeded changes
1 parent 5bf6e18 commit 92a0710

File tree

2 files changed

+44
-81
lines changed

2 files changed

+44
-81
lines changed

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

Lines changed: 26 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -420,10 +420,8 @@ describe('AmazonQTokenServiceManager', () => {
420420
const service = amazonQTokenServiceManager.getCodewhispererService()
421421
assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED')
422422
assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter')
423-
assert(codewhispererStubFactory.calledOnceWithExactly('eu-central-1', TEST_ENDPOINT_EU_CENTRAL_1))
424423

425-
assert.strictEqual(results[0].status, 'rejected')
426-
assert(results[0].reason.message, 'Requested profile update got cancelled')
424+
assert.strictEqual(results[0].status, 'fulfilled')
427425
assert.strictEqual(results[1].status, 'fulfilled')
428426
})
429427

@@ -552,7 +550,8 @@ describe('AmazonQTokenServiceManager', () => {
552550
assert.strictEqual(await streamingClient2.client.config.region(), 'eu-central-1')
553551
})
554552

555-
it('handles Profile configuration change from valid to invalid profile', async () => {
553+
// As we're not validating profile at this moment, there is no "invalid" profile
554+
it.skip('handles Profile configuration change from valid to invalid profile', async () => {
556555
setupServiceManager(true)
557556
assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION')
558557

@@ -616,7 +615,8 @@ describe('AmazonQTokenServiceManager', () => {
616615
assert.deepStrictEqual(codewhispererStubFactory.lastCall.args, ['us-east-1', TEST_ENDPOINT_US_EAST_1])
617616
})
618617

619-
it('handles non-existing profile selection', async () => {
618+
// As we're not validating profile at this moment, there is no "non-existing" profile
619+
it.skip('handles non-existing profile selection', async () => {
620620
setupServiceManager(true)
621621
assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION')
622622

@@ -663,25 +663,9 @@ describe('AmazonQTokenServiceManager', () => {
663663
)
664664
assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE')
665665

666-
await amazonQTokenServiceManager.handleOnUpdateConfiguration(
667-
{
668-
section: 'aws.q',
669-
settings: {
670-
profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ',
671-
},
672-
},
673-
{} as CancellationToken
674-
)
666+
amazonQTokenServiceManager.setState('PENDING_Q_PROFILE_UPDATE')
667+
assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE')
675668

676-
const pendingProfileUpdate = amazonQTokenServiceManager.handleOnUpdateConfiguration(
677-
{
678-
section: 'aws.q',
679-
settings: {
680-
profileArn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ',
681-
},
682-
},
683-
{} as CancellationToken
684-
)
685669
assert.throws(
686670
() => amazonQTokenServiceManager.getCodewhispererService(),
687671
AmazonQServicePendingProfileUpdateError
@@ -691,9 +675,15 @@ describe('AmazonQTokenServiceManager', () => {
691675
AmazonQServicePendingProfileUpdateError
692676
)
693677

694-
assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE')
695-
696-
await pendingProfileUpdate
678+
await amazonQTokenServiceManager.handleOnUpdateConfiguration(
679+
{
680+
section: 'aws.q',
681+
settings: {
682+
profileArn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ',
683+
},
684+
},
685+
{} as CancellationToken
686+
)
697687

698688
const service = amazonQTokenServiceManager.getCodewhispererService()
699689
const streamingClient = amazonQTokenServiceManager.getStreamingClient()
@@ -752,15 +742,7 @@ describe('AmazonQTokenServiceManager', () => {
752742
assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1')
753743

754744
// Updaing profile
755-
const pendingProfileUpdate = amazonQTokenServiceManager.handleOnUpdateConfiguration(
756-
{
757-
section: 'aws.q',
758-
settings: {
759-
profileArn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ',
760-
},
761-
},
762-
{} as CancellationToken
763-
)
745+
amazonQTokenServiceManager.setState('PENDING_Q_PROFILE_UPDATE')
764746
assert.throws(
765747
() => amazonQTokenServiceManager.getCodewhispererService(),
766748
AmazonQServicePendingProfileUpdateError
@@ -771,8 +753,6 @@ describe('AmazonQTokenServiceManager', () => {
771753
)
772754

773755
assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE')
774-
775-
await pendingProfileUpdate
776756
})
777757

778758
it('resets to PENDING_PROFILE from INITIALIZED when receiving null profileArn', async () => {
@@ -796,19 +776,12 @@ describe('AmazonQTokenServiceManager', () => {
796776
it('resets to PENDING_Q_PROFILE from PENDING_Q_PROFILE_UPDATE when receiving null profileArn', async () => {
797777
await setupServiceManagerWithProfile()
798778

799-
const pendingUpdate = amazonQTokenServiceManager.handleOnUpdateConfiguration(
800-
{
801-
section: 'aws.q',
802-
settings: {
803-
profileArn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ',
804-
},
805-
},
806-
{} as CancellationToken
807-
)
779+
amazonQTokenServiceManager.setState('PENDING_Q_PROFILE_UPDATE')
808780

809781
assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE')
810782

811-
const nullRequest = amazonQTokenServiceManager.handleOnUpdateConfiguration(
783+
// Null profile arn
784+
await amazonQTokenServiceManager.handleOnUpdateConfiguration(
812785
{
813786
section: 'aws.q',
814787
settings: {
@@ -818,8 +791,6 @@ describe('AmazonQTokenServiceManager', () => {
818791
{} as CancellationToken
819792
)
820793

821-
await Promise.allSettled([pendingUpdate, nullRequest])
822-
823794
assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE')
824795
assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined)
825796
sinon.assert.calledOnce(codewhispererServiceStub.abortInflightRequests)
@@ -829,31 +800,22 @@ describe('AmazonQTokenServiceManager', () => {
829800
it('cancels on-going profile update when credentials are deleted', async () => {
830801
await setupServiceManagerWithProfile()
831802

832-
const pendingUpdate = amazonQTokenServiceManager.handleOnUpdateConfiguration(
833-
{
834-
section: 'aws.q',
835-
settings: {
836-
profileArn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ',
837-
},
838-
},
839-
{} as CancellationToken
840-
)
841-
803+
amazonQTokenServiceManager.setState('PENDING_Q_PROFILE_UPDATE')
842804
assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE')
843805

844806
amazonQTokenServiceManager.handleOnCredentialsDeleted('bearer')
845807

846808
assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION')
847809

848-
await assert.rejects(() => pendingUpdate)
849-
850810
assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION')
851811
assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined)
852812
sinon.assert.calledOnce(codewhispererServiceStub.abortInflightRequests)
853813
assert.throws(() => amazonQTokenServiceManager.getCodewhispererService())
854814
})
855815

856-
it('fetches profiles only from 1 region associated with requested profileArn', async () => {
816+
// Due to service limitation, validation was removed for the sake of recovering API availability
817+
// When service is ready to take more tps, revert https://github.com/aws/language-servers/pull/1329 to add profile validation
818+
it('should not call service to validate profile and always assume its validness', async () => {
857819
setupServiceManager(true)
858820
assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION')
859821

@@ -869,9 +831,8 @@ describe('AmazonQTokenServiceManager', () => {
869831
{} as CancellationToken
870832
)
871833

872-
sinon.assert.calledOnceWithMatch(getListAllAvailableProfilesHandlerStub, {
873-
endpoints: new Map([['us-east-1', TEST_ENDPOINT_US_EAST_1]]),
874-
})
834+
sinon.assert.notCalled(getListAllAvailableProfilesHandlerStub)
835+
assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED')
875836
})
876837
})
877838
})

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

Lines changed: 18 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,7 @@ import {
2727
QServiceManagerFeatures,
2828
} from './BaseAmazonQServiceManager'
2929
import { AWS_Q_ENDPOINTS, Q_CONFIGURATION_SECTION } from '../constants'
30-
import {
31-
AmazonQDeveloperProfile,
32-
getListAllAvailableProfilesHandler,
33-
signalsAWSQDeveloperProfilesEnabled,
34-
} from './qDeveloperProfiles'
30+
import { AmazonQDeveloperProfile, signalsAWSQDeveloperProfilesEnabled } from './qDeveloperProfiles'
3531
import { isStringOrNull } from '../utils'
3632
import { getAmazonQRegionAndEndpoint } from './configurationUtils'
3733
import { getUserAgent } from '../telemetryUtils'
@@ -89,6 +85,11 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
8985
super(features)
9086
}
9187

88+
// @VisibleForTesting, please DO NOT use in production
89+
setState(state: 'PENDING_CONNECTION' | 'PENDING_Q_PROFILE' | 'PENDING_Q_PROFILE_UPDATE' | 'INITIALIZED') {
90+
this.state = state
91+
}
92+
9293
public static initInstance(features: QServiceManagerFeatures): AmazonQTokenServiceManager {
9394
if (!AmazonQTokenServiceManager.instance) {
9495
AmazonQTokenServiceManager.instance = new AmazonQTokenServiceManager(features)
@@ -151,6 +152,7 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
151152
try {
152153
if (params.section === Q_CONFIGURATION_SECTION && params.settings.profileArn !== undefined) {
153154
const profileArn = params.settings.profileArn
155+
const region = params.settings.region
154156

155157
if (!isStringOrNull(profileArn)) {
156158
throw new Error('Expected params.settings.profileArn to be of either type string or null')
@@ -318,21 +320,21 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<
318320
}
319321

320322
const parsedArn = parse(newProfileArn)
321-
const endpoint = AWS_Q_ENDPOINTS.get(parsedArn.region)
323+
const region = parsedArn.region
324+
const endpoint = AWS_Q_ENDPOINTS.get(region)
322325
if (!endpoint) {
323326
throw new Error('Requested profileArn region is not supported')
324327
}
325328

326-
const profiles = await getListAllAvailableProfilesHandler(this.serviceFactory)({
327-
connectionType: 'identityCenter',
328-
logging: this.logging,
329-
token: token,
330-
endpoints: new Map([[parsedArn.region, endpoint]]),
331-
})
332-
333-
this.handleTokenCancellationRequest(token)
334-
335-
const newProfile = profiles.find(el => el.arn === newProfileArn)
329+
// Hack to inject a dummy profile name as it's not used by client IDE for now, if client IDE starts consuming name field then we should also pass both profile name and arn from the IDE
330+
// When service is ready to take more tps, revert https://github.com/aws/language-servers/pull/1329 to add profile validation
331+
const newProfile: AmazonQDeveloperProfile = {
332+
arn: newProfileArn,
333+
name: 'Client provided profile',
334+
identityDetails: {
335+
region: parsedArn.region,
336+
},
337+
}
336338

337339
if (!newProfile || !newProfile.identityDetails?.region) {
338340
this.log(`Amazon Q Profile ${newProfileArn} is not valid`)

0 commit comments

Comments
 (0)