Skip to content

Commit ff0caa4

Browse files
Will-ShaoHuaWill-ShaoHua
committed
show all customizations across different profiles for users to choose
1 parent d6d839a commit ff0caa4

File tree

6 files changed

+96
-60
lines changed

6 files changed

+96
-60
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -240,7 +240,7 @@ describe('RegionProfileManager', function () {
240240
assert.deepStrictEqual(sut.activeRegionProfile, profileFoo)
241241
const conn = authUtil.conn as SsoConnection
242242

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

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

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import { indent } from '../../shared/utilities/textUtilities'
2323
import { getClientId, getOptOutPreference, getOperatingSystem } from '../../shared/telemetry/util'
2424
import { extensionVersion, getServiceEnvVarConfig } from '../../shared/vscode/env'
2525
import { DevSettings } from '../../shared/settings'
26-
import { CodeWhispererConfig } from '../models/model'
26+
import { CodeWhispererConfig, RegionProfile } from '../models/model'
2727

2828
const keepAliveHeader = 'keep-alive-codewhisperer'
2929

@@ -219,9 +219,8 @@ export class DefaultCodeWhispererClient {
219219
.promise()
220220
}
221221

222-
public async listAvailableCustomizations(): Promise<ListAvailableCustomizationsResponse[]> {
223-
const client = await this.createUserSdkClient()
224-
const profile = AuthUtil.instance.regionProfileManager.activeRegionProfile
222+
public async listAvailableCustomizations(profile: RegionProfile): Promise<ListAvailableCustomizationsResponse[]> {
223+
const client = await AuthUtil.instance.regionProfileManager.createQClient(profile)
225224
const requester = async (request: CodeWhispererUserClient.ListAvailableCustomizationsRequest) =>
226225
client.listAvailableCustomizations(request).promise()
227226
return pageableToCollection(requester, { profileArn: profile?.arn }, 'nextToken')

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

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ const endpoints = createConstantMap({
4747
* 'update' -> plugin auto select the profile on users' behalf as there is only 1 profile
4848
* 'reload' -> on plugin restart, plugin will try to reload previous selected profile
4949
*/
50-
export type ProfileSwitchIntent = 'user' | 'auth' | 'update' | 'reload'
50+
export type ProfileSwitchIntent = 'user' | 'auth' | 'update' | 'reload' | 'customization'
5151

5252
export class RegionProfileManager {
5353
private static logger = getLogger()
@@ -112,7 +112,7 @@ export class RegionProfileManager {
112112
}
113113
const availableProfiles: RegionProfile[] = []
114114
for (const [region, endpoint] of endpoints.entries()) {
115-
const client = await this.createQClient(region, endpoint, conn as SsoConnection)
115+
const client = await this._createQClient(region, endpoint, conn as SsoConnection)
116116
const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) =>
117117
client.listAvailableProfiles(request).promise()
118118
const request: CodeWhispererUserClient.ListAvailableProfilesRequest = {}
@@ -162,7 +162,8 @@ export class RegionProfileManager {
162162
const ssoConn = this.connectionProvider() as SsoConnection
163163

164164
// only prompt to users when users switch from A profile to B profile
165-
if (this.activeRegionProfile !== undefined && regionProfile !== undefined) {
165+
// TODO: should we ask if it's a customization switch?
166+
if (source !== 'customization' && this.activeRegionProfile !== undefined && regionProfile !== undefined) {
166167
const response = await showConfirmationMessage({
167168
prompt: localize(
168169
'AWS.amazonq.profile.confirmation',
@@ -343,7 +344,20 @@ export class RegionProfileManager {
343344
}
344345
}
345346

346-
async createQClient(region: string, endpoint: string, conn: SsoConnection): Promise<CodeWhispererUserClient> {
347+
// TODO: overload?
348+
async createQClient(profile: RegionProfile): Promise<CodeWhispererUserClient> {
349+
const conn = this.connectionProvider()
350+
if (conn === undefined || !isSsoConnection(conn)) {
351+
throw new Error('No valid SSO connection')
352+
}
353+
const endpoint = endpoints.get(profile.region)
354+
if (!endpoint) {
355+
throw new Error(`trying to initiatize Q client with unrecognizable region ${profile.region}`)
356+
}
357+
return this._createQClient(profile.region, endpoint, conn)
358+
}
359+
360+
async _createQClient(region: string, endpoint: string, conn: SsoConnection): Promise<CodeWhispererUserClient> {
347361
const token = (await conn.getToken()).accessToken
348362
const serviceOption: ServiceOptions = {
349363
apiConfig: userApiConfig,

‎packages/core/src/codewhisperer/util/customizationUtil.ts‎

Lines changed: 25 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import { getLogger } from '../../shared/logger/logger'
1717
import { showMessageWithUrl } from '../../shared/utilities/messages'
1818
import { parse } from '@aws-sdk/util-arn-parser'
1919
import { Commands } from '../../shared/vscode/commands2'
20-
import { vsCodeState } from '../models/model'
20+
import { RegionProfile, vsCodeState } from '../models/model'
2121

2222
/**
2323
*
@@ -224,7 +224,6 @@ const createCustomizationItems = async () => {
224224
if (availableCustomizations.length === 0) {
225225
items.push(createBaseCustomizationItem())
226226

227-
// TODO: finalize the url string with documentation
228227
void showMessageWithUrl(
229228
localize(
230229
'AWS.codewhisperer.customization.noCustomizations.description',
@@ -284,7 +283,7 @@ const createBaseCustomizationItem = () => {
284283
}
285284

286285
const createCustomizationItem = (
287-
customization: Customization,
286+
customization: Customization & { profile: RegionProfile },
288287
persistedArns: (ResourceArn | undefined)[],
289288
shouldPrefixAccountId: boolean
290289
) => {
@@ -293,8 +292,8 @@ const createCustomizationItem = (
293292
? shouldPrefixAccountId
294293
? accountId
295294
? `${customization.name} (${accountId})`
296-
: `${customization.name}`
297-
: customization.name
295+
: `${customization.name} (${customization.profile.name})`
296+
: `${customization.name} (${customization.profile.name})`
298297
: 'unknown'
299298

300299
const isNewCustomization = !persistedArns.includes(customization.arn)
@@ -303,6 +302,10 @@ const createCustomizationItem = (
303302
return {
304303
label: label,
305304
onClick: async () => {
305+
const profile = AuthUtil.instance.regionProfileManager.activeRegionProfile
306+
if (profile && customization.profile.arn !== profile.arn) {
307+
await AuthUtil.instance.regionProfileManager.switchRegionProfile(customization.profile, 'customization')
308+
}
306309
await selectCustomization(customization)
307310
},
308311
detail:
@@ -333,13 +336,24 @@ export const selectCustomization = async (customization: Customization) => {
333336
)
334337
}
335338

339+
// should collect all customizations across different profiles and if users select the customization, we also change the profile if needed if the customization is accessible from a different profile
336340
export const getAvailableCustomizationsList = async () => {
337-
const items: Customization[] = []
338-
const response = await codeWhispererClient.listAvailableCustomizations()
339-
for (const customizations of response.map(
340-
(listAvailableCustomizationsResponse) => listAvailableCustomizationsResponse.customizations
341-
)) {
342-
items.push(...customizations)
341+
const items: (Customization & { profile: RegionProfile })[] = []
342+
// TODO: try catch listProfile
343+
const profiles = await AuthUtil.instance.regionProfileManager.listRegionProfile()
344+
345+
for (const profile of profiles) {
346+
const response = await codeWhispererClient.listAvailableCustomizations(profile)
347+
for (const customizations of response.map(
348+
(listAvailableCustomizationsResponse) => listAvailableCustomizationsResponse.customizations
349+
)) {
350+
for (const c of customizations) {
351+
items.push({
352+
...c,
353+
profile: profile,
354+
})
355+
}
356+
}
343357
}
344358

345359
return items

‎packages/core/src/shared/featureConfig.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -145,18 +145,20 @@ export class FeatureConfigProvider {
145145
}
146146
getLogger().info('AB Testing Cohort Assignments %O', response.featureEvaluations)
147147

148+
const profile = AuthUtil.instance.regionProfileManager.activeRegionProfile
148149
const customizationArnOverride = this.featureConfigs.get(Features.customizationArnOverride)?.value
149150
?.stringValue
150151
const previousOverride = globals.globalState.tryGet<string>('aws.amazonq.customization.overrideV2', String)
151-
if (customizationArnOverride !== undefined && customizationArnOverride !== previousOverride) {
152+
if (profile && customizationArnOverride !== undefined && customizationArnOverride !== previousOverride) {
152153
// Double check if server-side wrongly returns a customizationArn to BID users
153154
if (isBuilderIdConnection(AuthUtil.instance.conn)) {
154155
this.featureConfigs.delete(Features.customizationArnOverride)
155156
} else if (isIdcSsoConnection(AuthUtil.instance.conn)) {
156157
let availableCustomizations: Customization[] = []
157158
try {
158159
const items: Customization[] = []
159-
const response = await client.listAvailableCustomizations()
160+
// TODO: list all customization across different profiles
161+
const response = await client.listAvailableCustomizations(profile)
160162
for (const customizations of response.map(
161163
(listAvailableCustomizationsResponse) => listAvailableCustomizationsResponse.customizations
162164
)) {

0 commit comments

Comments
 (0)