Skip to content

Commit f30f770

Browse files
Will-ShaoHuaWill-ShaoHua
authored
feat(amazonq): show all customizations across different profiles aws#7181
this effectively restores aws#7060 (and reverts "revert PR" aws#7129)
1 parent 1fb7013 commit f30f770

File tree

9 files changed

+249
-90
lines changed

9 files changed

+249
-90
lines changed

‎packages/amazonq/.changes/next-release/Feature-2a696d00-e8c8-44a4-ab2c-2204b4d8e31d.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "Feature",
3+
"description": "Support selecting customizations across all Q profiles with automatic profile switching for enterprise users"
4+
}

‎packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts‎

Lines changed: 54 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ describe('RegionProfileManager', function () {
6565
const mockClient = {
6666
listAvailableProfiles: listProfilesStub,
6767
}
68-
const createClientStub = sinon.stub(sut, 'createQClient').resolves(mockClient)
68+
const createClientStub = sinon.stub(sut, '_createQClient').resolves(mockClient)
6969

7070
const r = await sut.listRegionProfile()
7171

@@ -234,13 +234,65 @@ describe('RegionProfileManager', function () {
234234
})
235235

236236
describe('createQClient', function () {
237+
it(`should configure the endpoint and region from a profile`, async function () {
238+
await setupConnection('idc')
239+
240+
const iadClient = await sut.createQClient({
241+
name: 'foo',
242+
region: 'us-east-1',
243+
arn: 'arn',
244+
description: 'description',
245+
})
246+
247+
assert.deepStrictEqual(iadClient.config.region, 'us-east-1')
248+
assert.deepStrictEqual(iadClient.endpoint.href, 'https://q.us-east-1.amazonaws.com/')
249+
250+
const fraClient = await sut.createQClient({
251+
name: 'bar',
252+
region: 'eu-central-1',
253+
arn: 'arn',
254+
description: 'description',
255+
})
256+
257+
assert.deepStrictEqual(fraClient.config.region, 'eu-central-1')
258+
assert.deepStrictEqual(fraClient.endpoint.href, 'https://q.eu-central-1.amazonaws.com/')
259+
})
260+
261+
it(`should throw if the region is not supported or recognizable by Q`, async function () {
262+
await setupConnection('idc')
263+
264+
await assert.rejects(
265+
async () => {
266+
await sut.createQClient({
267+
name: 'foo',
268+
region: 'ap-east-1',
269+
arn: 'arn',
270+
description: 'description',
271+
})
272+
},
273+
{ message: /trying to initiatize Q client with unrecognizable region/ }
274+
)
275+
276+
await assert.rejects(
277+
async () => {
278+
await sut.createQClient({
279+
name: 'foo',
280+
region: 'unknown-somewhere',
281+
arn: 'arn',
282+
description: 'description',
283+
})
284+
},
285+
{ message: /trying to initiatize Q client with unrecognizable region/ }
286+
)
287+
})
288+
237289
it(`should configure the endpoint and region correspondingly`, async function () {
238290
await setupConnection('idc')
239291
await sut.switchRegionProfile(profileFoo, 'user')
240292
assert.deepStrictEqual(sut.activeRegionProfile, profileFoo)
241293
const conn = authUtil.conn as SsoConnection
242294

243-
const client = await sut.createQClient('eu-central-1', 'https://amazon.com/', conn)
295+
const client = await sut._createQClient('eu-central-1', 'https://amazon.com/', conn)
244296

245297
assert.deepStrictEqual(client.config.region, 'eu-central-1')
246298
assert.deepStrictEqual(client.endpoint.href, 'https://amazon.com/')

‎packages/core/src/codewhisperer/activation.ts‎

Lines changed: 2 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -72,12 +72,7 @@ import { AuthUtil } from './util/authUtil'
7272
import { ImportAdderProvider } from './service/importAdderProvider'
7373
import { TelemetryHelper } from './util/telemetryHelper'
7474
import { openUrl } from '../shared/utilities/vsCodeUtils'
75-
import {
76-
getAvailableCustomizationsList,
77-
getSelectedCustomization,
78-
notifyNewCustomizations,
79-
switchToBaseCustomizationAndNotify,
80-
} from './util/customizationUtil'
75+
import { notifyNewCustomizations, onProfileChangedListener } from './util/customizationUtil'
8176
import { CodeWhispererCommandBackend, CodeWhispererCommandDeclarations } from './commands/gettingStartedPageCommands'
8277
import { SecurityIssueHoverProvider } from './service/securityIssueHoverProvider'
8378
import { SecurityIssueCodeActionProvider } from './service/securityIssueCodeActionProvider'
@@ -344,26 +339,7 @@ export async function activate(context: ExtContext): Promise<void> {
344339
SecurityIssueCodeActionProvider.instance
345340
),
346341
vscode.commands.registerCommand('aws.amazonq.openEditorAtRange', openEditorAtRange),
347-
auth.regionProfileManager.onDidChangeRegionProfile(() => {
348-
// Validate user still has access to the selected customization.
349-
const selectedCustomization = getSelectedCustomization()
350-
// No need to validate base customization which has empty arn.
351-
if (selectedCustomization.arn.length > 0) {
352-
getAvailableCustomizationsList()
353-
.then(async (customizations) => {
354-
const r = customizations.find((it) => it.arn === selectedCustomization.arn)
355-
if (!r) {
356-
await switchToBaseCustomizationAndNotify()
357-
}
358-
})
359-
.catch((e) => {
360-
getLogger().error(
361-
`encounter error while validating selected customization on profile change: %s`,
362-
(e as Error).message
363-
)
364-
})
365-
}
366-
})
342+
auth.regionProfileManager.onDidChangeRegionProfile(onProfileChangedListener)
367343
)
368344

369345
// run the auth startup code with context for telemetry

‎packages/core/src/codewhisperer/client/codewhisperer.ts‎

Lines changed: 1 addition & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -7,19 +7,17 @@ import { AWSError, Credentials, Service } from 'aws-sdk'
77
import globals from '../../shared/extensionGlobals'
88
import * as CodeWhispererClient from './codewhispererclient'
99
import * as CodeWhispererUserClient from './codewhispereruserclient'
10-
import { ListAvailableCustomizationsResponse, SendTelemetryEventRequest } from './codewhispereruserclient'
10+
import { SendTelemetryEventRequest } from './codewhispereruserclient'
1111
import { ServiceOptions } from '../../shared/awsClientBuilder'
1212
import { hasVendedIamCredentials } from '../../auth/auth'
1313
import { CodeWhispererSettings } from '../util/codewhispererSettings'
1414
import { PromiseResult } from 'aws-sdk/lib/request'
1515
import { AuthUtil } from '../util/authUtil'
1616
import { isSsoConnection } from '../../auth/connection'
17-
import { pageableToCollection } from '../../shared/utilities/collectionUtils'
1817
import apiConfig = require('./service-2.json')
1918
import userApiConfig = require('./user-service-2.json')
2019
import { session } from '../util/codeWhispererSession'
2120
import { getLogger } from '../../shared/logger/logger'
22-
import { indent } from '../../shared/utilities/textUtilities'
2321
import { getClientId, getOptOutPreference, getOperatingSystem } from '../../shared/telemetry/util'
2422
import { extensionVersion, getServiceEnvVarConfig } from '../../shared/vscode/env'
2523
import { DevSettings } from '../../shared/settings'
@@ -219,28 +217,6 @@ export class DefaultCodeWhispererClient {
219217
.promise()
220218
}
221219

222-
public async listAvailableCustomizations(): Promise<ListAvailableCustomizationsResponse[]> {
223-
const client = await this.createUserSdkClient()
224-
const profile = AuthUtil.instance.regionProfileManager.activeRegionProfile
225-
const requester = async (request: CodeWhispererUserClient.ListAvailableCustomizationsRequest) =>
226-
client.listAvailableCustomizations(request).promise()
227-
return pageableToCollection(requester, { profileArn: profile?.arn }, 'nextToken')
228-
.promise()
229-
.then((resps) => {
230-
let logStr = 'amazonq: listAvailableCustomizations API request:'
231-
for (const resp of resps) {
232-
const requestId = resp.$response.requestId
233-
logStr += `\n${indent('RequestID: ', 4)}${requestId},\n${indent('Customizations:', 4)}`
234-
for (const [index, c] of resp.customizations.entries()) {
235-
const entry = `${index.toString().padStart(2, '0')}: ${c.name?.trim()}`
236-
logStr += `\n${indent(entry, 8)}`
237-
}
238-
}
239-
getLogger().debug(logStr)
240-
return resps
241-
})
242-
}
243-
244220
public async sendTelemetryEvent(request: SendTelemetryEventRequest) {
245221
const requestWithCommonFields: SendTelemetryEventRequest = {
246222
...request,

‎packages/core/src/codewhisperer/index.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,13 @@ export * as diagnosticsProvider from './service/diagnosticsProvider'
9999
export * from './ui/codeWhispererNodes'
100100
export { SecurityScanError, SecurityScanTimedOutError } from '../codewhisperer/models/errors'
101101
export * as CodeWhispererConstants from '../codewhisperer/models/constants'
102-
export { getSelectedCustomization, setSelectedCustomization, baseCustomization } from './util/customizationUtil'
102+
export {
103+
getSelectedCustomization,
104+
setSelectedCustomization,
105+
baseCustomization,
106+
onProfileChangedListener,
107+
CustomizationProvider,
108+
} from './util/customizationUtil'
103109
export { Container } from './service/serviceContainer'
104110
export * from './util/gitUtil'
105111
export * from './ui/prompters'

‎packages/core/src/codewhisperer/region/regionProfileManager.ts‎

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -49,12 +49,17 @@ const endpoints = createConstantMap({
4949
* 'update' -> plugin auto select the profile on users' behalf as there is only 1 profile
5050
* 'reload' -> on plugin restart, plugin will try to reload previous selected profile
5151
*/
52-
export type ProfileSwitchIntent = 'user' | 'auth' | 'update' | 'reload'
52+
export type ProfileSwitchIntent = 'user' | 'auth' | 'update' | 'reload' | 'customization'
53+
54+
export type ProfileChangedEvent = {
55+
profile: RegionProfile | undefined
56+
intent: ProfileSwitchIntent
57+
}
5358

5459
export class RegionProfileManager {
5560
private static logger = getLogger()
5661
private _activeRegionProfile: RegionProfile | undefined
57-
private _onDidChangeRegionProfile = new vscode.EventEmitter<RegionProfile | undefined>()
62+
private _onDidChangeRegionProfile = new vscode.EventEmitter<ProfileChangedEvent>()
5863
public readonly onDidChangeRegionProfile = this._onDidChangeRegionProfile.event
5964

6065
// Store the last API results (for UI propuse) so we don't need to call service again if doesn't require "latest" result
@@ -141,7 +146,7 @@ export class RegionProfileManager {
141146
const failedRegions: string[] = []
142147

143148
for (const [region, endpoint] of endpoints.entries()) {
144-
const client = await this.createQClient(region, endpoint, conn as SsoConnection)
149+
const client = await this._createQClient(region, endpoint, conn as SsoConnection)
145150
const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) =>
146151
client.listAvailableProfiles(request).promise()
147152
const request: CodeWhispererUserClient.ListAvailableProfilesRequest = {}
@@ -195,7 +200,7 @@ export class RegionProfileManager {
195200
const ssoConn = this.connectionProvider() as SsoConnection
196201

197202
// only prompt to users when users switch from A profile to B profile
198-
if (this.activeRegionProfile !== undefined && regionProfile !== undefined) {
203+
if (source !== 'customization' && this.activeRegionProfile !== undefined && regionProfile !== undefined) {
199204
const response = await showConfirmationMessage({
200205
prompt: localize(
201206
'AWS.amazonq.profile.confirmation',
@@ -237,13 +242,16 @@ export class RegionProfileManager {
237242
})
238243
}
239244

240-
await this._switchRegionProfile(regionProfile)
245+
await this._switchRegionProfile(regionProfile, source)
241246
}
242247

243-
private async _switchRegionProfile(regionProfile: RegionProfile | undefined) {
248+
private async _switchRegionProfile(regionProfile: RegionProfile | undefined, source: ProfileSwitchIntent) {
244249
this._activeRegionProfile = regionProfile
245250

246-
this._onDidChangeRegionProfile.fire(regionProfile)
251+
this._onDidChangeRegionProfile.fire({
252+
profile: regionProfile,
253+
intent: source,
254+
})
247255
// dont show if it's a default (fallback)
248256
if (regionProfile && this.profiles.length > 1) {
249257
void vscode.window.showInformationMessage(`You are using the ${regionProfile.name} profile for Q.`).then()
@@ -384,7 +392,21 @@ export class RegionProfileManager {
384392
await this.cache.clearCache()
385393
}
386394

387-
async createQClient(region: string, endpoint: string, conn: SsoConnection): Promise<CodeWhispererUserClient> {
395+
// TODO: Should maintain sdk client in a better way
396+
async createQClient(profile: RegionProfile): Promise<CodeWhispererUserClient> {
397+
const conn = this.connectionProvider()
398+
if (conn === undefined || !isSsoConnection(conn)) {
399+
throw new Error('No valid SSO connection')
400+
}
401+
const endpoint = endpoints.get(profile.region)
402+
if (!endpoint) {
403+
throw new Error(`trying to initiatize Q client with unrecognizable region ${profile.region}`)
404+
}
405+
return this._createQClient(profile.region, endpoint, conn)
406+
}
407+
408+
// Visible for testing only, do not use this directly, please use createQClient(profile)
409+
async _createQClient(region: string, endpoint: string, conn: SsoConnection): Promise<CodeWhispererUserClient> {
388410
const token = (await conn.getToken()).accessToken
389411
const serviceOption: ServiceOptions = {
390412
apiConfig: userApiConfig,

0 commit comments

Comments
 (0)