Skip to content

Commit df8cb90

Browse files
Will-ShaoHuaWill-ShaoHua
committed
Merge remote-tracking branch 'upstream/feature/q-region-expansion' into chore
2 parents 191c504 + bab1a70 commit df8cb90

File tree

66 files changed

+3089
-151
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

66 files changed

+3089
-151
lines changed

‎packages/amazonq/package.json‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -357,6 +357,11 @@
357357
"when": "view =~ /^aws\\.amazonq/",
358358
"group": "1_amazonQ@1"
359359
},
360+
{
361+
"command": "aws.amazonq.selectRegionProfile",
362+
"when": "view == aws.amazonq.AmazonQChatView && aws.amazonq.connectedSsoIdc == true",
363+
"group": "1_amazonQ@1"
364+
},
360365
{
361366
"command": "aws.amazonq.signout",
362367
"when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected && !aws.isSageMakerUnifiedStudio",
@@ -567,6 +572,12 @@
567572
"category": "%AWS.amazonq.title%",
568573
"enablement": "aws.codewhisperer.connected"
569574
},
575+
{
576+
"command": "aws.amazonq.selectRegionProfile",
577+
"title": "Change Profile",
578+
"category": "%AWS.amazonq.title%",
579+
"enablement": "aws.codewhisperer.connected"
580+
},
570581
{
571582
"command": "aws.amazonq.transformationHub.reviewChanges.acceptChanges",
572583
"title": "%AWS.command.q.transform.acceptChanges%"

‎packages/amazonq/src/app/amazonqScan/app.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,9 @@ export function init(appContext: AmazonQAppInitContext) {
7070
AuthUtil.instance.secondaryAuth.onDidChangeActiveConnection(() => {
7171
return debouncedEvent()
7272
})
73+
AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(() => {
74+
return debouncedEvent()
75+
})
7376

7477
Commands.register('aws.amazonq.security.scan-statusbar', async () => {
7578
if (AuthUtil.instance.isConnectionExpired()) {

‎packages/amazonq/src/app/amazonqScan/chat/controller/controller.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,10 @@ export class ScanController {
108108
interactionType: data.vote,
109109
})
110110
})
111+
112+
AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(() => {
113+
this.sessionStorage.removeActiveTab()
114+
})
111115
}
112116

113117
private async tabOpened(message: any) {

‎packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ export class InlineChatProvider {
6666
codeQuery: context?.focusAreaContext?.names,
6767
userIntent: this.userIntentRecognizer.getFromPromptChatMessage(message),
6868
customization: getSelectedCustomization(),
69+
profile: AuthUtil.instance.regionProfileManager.activeRegionProfile,
6970
context: [],
7071
relevantTextDocuments: [],
7172
additionalContents: [],

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

Lines changed: 249 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,249 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import * as sinon from 'sinon'
7+
import assert, { fail } from 'assert'
8+
import { AuthUtil, RegionProfile, RegionProfileManager, defaultServiceConfig } from 'aws-core-vscode/codewhisperer'
9+
import { globals } from 'aws-core-vscode/shared'
10+
import { createTestAuth } from 'aws-core-vscode/test'
11+
import { SsoConnection } from 'aws-core-vscode/auth'
12+
13+
const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start'
14+
15+
describe('RegionProfileManager', function () {
16+
let sut: RegionProfileManager
17+
let auth: ReturnType<typeof createTestAuth>
18+
let authUtil: AuthUtil
19+
20+
const profileFoo: RegionProfile = {
21+
name: 'foo',
22+
region: 'us-east-1',
23+
arn: 'foo arn',
24+
description: 'foo description',
25+
}
26+
27+
async function setupConnection(type: 'builderId' | 'idc') {
28+
if (type === 'builderId') {
29+
await authUtil.connectToAwsBuilderId()
30+
const conn = authUtil.conn
31+
assert.strictEqual(conn?.type, 'sso')
32+
assert.strictEqual(conn.label, 'AWS Builder ID')
33+
} else if (type === 'idc') {
34+
await authUtil.connectToEnterpriseSso(enterpriseSsoStartUrl, 'us-east-1')
35+
const conn = authUtil.conn
36+
assert.strictEqual(conn?.type, 'sso')
37+
assert.strictEqual(conn.label, 'IAM Identity Center (enterprise)')
38+
}
39+
}
40+
41+
beforeEach(function () {
42+
auth = createTestAuth(globals.globalState)
43+
authUtil = new AuthUtil(auth)
44+
sut = new RegionProfileManager(() => authUtil.conn)
45+
})
46+
47+
afterEach(function () {
48+
sinon.restore()
49+
})
50+
51+
describe('list profiles', function () {
52+
it('should call list profiles with different region endpoints', async function () {
53+
await setupConnection('idc')
54+
const listProfilesStub = sinon.stub().returns({
55+
promise: () =>
56+
Promise.resolve({
57+
profiles: [
58+
{
59+
arn: 'arn',
60+
profileName: 'foo',
61+
},
62+
],
63+
}),
64+
})
65+
const mockClient = {
66+
listAvailableProfiles: listProfilesStub,
67+
}
68+
const createClientStub = sinon.stub(sut, 'createQClient').resolves(mockClient)
69+
70+
const r = await sut.listRegionProfile()
71+
72+
assert.strictEqual(r.length, 2)
73+
assert.deepStrictEqual(r, [
74+
{
75+
name: 'foo',
76+
arn: 'arn',
77+
region: 'us-east-1',
78+
description: '',
79+
},
80+
{
81+
name: 'foo',
82+
arn: 'arn',
83+
region: 'eu-central-1',
84+
description: '',
85+
},
86+
])
87+
88+
assert.ok(createClientStub.calledTwice)
89+
assert.ok(listProfilesStub.calledTwice)
90+
})
91+
})
92+
93+
describe('switch and get profile', function () {
94+
it('should switch if connection is IdC', async function () {
95+
await setupConnection('idc')
96+
await sut.switchRegionProfile(profileFoo, 'user')
97+
assert.deepStrictEqual(sut.activeRegionProfile, profileFoo)
98+
})
99+
100+
it('should do nothing and return undefined if connection is builder id', async function () {
101+
await setupConnection('builderId')
102+
await sut.switchRegionProfile(profileFoo, 'user')
103+
assert.deepStrictEqual(sut.activeRegionProfile, undefined)
104+
})
105+
})
106+
107+
describe(`client config`, function () {
108+
it(`no valid credential should throw`, async function () {
109+
assert.ok(authUtil.conn === undefined)
110+
111+
assert.throws(() => {
112+
sut.clientConfig
113+
}, /trying to get client configuration without credential/)
114+
})
115+
116+
it(`builder id should always use default profile IAD`, async function () {
117+
await setupConnection('builderId')
118+
await sut.switchRegionProfile(profileFoo, 'user')
119+
assert.deepStrictEqual(sut.activeRegionProfile, undefined)
120+
const conn = authUtil.conn
121+
if (!conn) {
122+
fail('connection should not be undefined')
123+
}
124+
125+
assert.deepStrictEqual(sut.clientConfig, defaultServiceConfig)
126+
})
127+
128+
it(`idc should return correct endpoint corresponding to profile region`, async function () {
129+
await setupConnection('idc')
130+
await sut.switchRegionProfile(
131+
{
132+
name: 'foo',
133+
region: 'eu-central-1',
134+
arn: 'foo arn',
135+
description: 'foo description',
136+
},
137+
'user'
138+
)
139+
assert.ok(sut.activeRegionProfile)
140+
assert.deepStrictEqual(sut.clientConfig, {
141+
region: 'eu-central-1',
142+
endpoint: 'https://q.eu-central-1.amazonaws.com/',
143+
})
144+
})
145+
146+
it(`idc should throw if corresponding endpoint is not defined`, async function () {
147+
await setupConnection('idc')
148+
await sut.switchRegionProfile(
149+
{
150+
name: 'foo',
151+
region: 'unknown region',
152+
arn: 'foo arn',
153+
description: 'foo description',
154+
},
155+
'user'
156+
)
157+
158+
assert.throws(() => {
159+
sut.clientConfig
160+
}, /Q client configuration error, endpoint not found for region*/)
161+
})
162+
})
163+
164+
describe('persistence', function () {
165+
it('persistSelectedRegionProfile', async function () {
166+
await setupConnection('idc')
167+
await sut.switchRegionProfile(profileFoo, 'user')
168+
assert.deepStrictEqual(sut.activeRegionProfile, profileFoo)
169+
const conn = authUtil.conn
170+
if (!conn) {
171+
fail('connection should not be undefined')
172+
}
173+
174+
await sut.persistSelectRegionProfile()
175+
176+
const state = globals.globalState.tryGet<{ [label: string]: RegionProfile }>(
177+
'aws.amazonq.regionProfiles',
178+
Object,
179+
{}
180+
)
181+
182+
assert.strictEqual(state[conn.id], profileFoo)
183+
})
184+
185+
it(`restoreRegionProfile`, async function () {
186+
sinon.stub(sut, 'listRegionProfile').resolves([profileFoo])
187+
await setupConnection('idc')
188+
const conn = authUtil.conn
189+
if (!conn) {
190+
fail('connection should not be undefined')
191+
}
192+
193+
const state = {} as any
194+
state[conn.id] = profileFoo
195+
196+
await globals.globalState.update('aws.amazonq.regionProfiles', state)
197+
198+
await sut.restoreRegionProfile(conn)
199+
200+
assert.strictEqual(sut.activeRegionProfile, profileFoo)
201+
})
202+
})
203+
204+
describe('invalidate', function () {
205+
it('should reset activeProfile and global state', async function () {
206+
// setup
207+
await setupConnection('idc')
208+
await sut.switchRegionProfile(profileFoo, 'user')
209+
assert.deepStrictEqual(sut.activeRegionProfile, profileFoo)
210+
const conn = authUtil.conn
211+
if (!conn) {
212+
fail('connection should not be undefined')
213+
}
214+
await sut.persistSelectRegionProfile()
215+
const state = globals.globalState.tryGet<{ [label: string]: RegionProfile }>(
216+
'aws.amazonq.regionProfiles',
217+
Object,
218+
{}
219+
)
220+
assert.strictEqual(state[conn.id], profileFoo)
221+
222+
// subject to test
223+
await sut.invalidateProfile(profileFoo.arn)
224+
225+
// assertion
226+
assert.strictEqual(sut.activeRegionProfile, undefined)
227+
const actualGlobalState = globals.globalState.tryGet<{ [label: string]: RegionProfile }>(
228+
'aws.amazonq.regionProfiles',
229+
Object,
230+
{}
231+
)
232+
assert.deepStrictEqual(actualGlobalState, {})
233+
})
234+
})
235+
236+
describe('createQClient', function () {
237+
it(`should configure the endpoint and region correspondingly`, async function () {
238+
await setupConnection('idc')
239+
await sut.switchRegionProfile(profileFoo, 'user')
240+
assert.deepStrictEqual(sut.activeRegionProfile, profileFoo)
241+
const conn = authUtil.conn as SsoConnection
242+
243+
const client = await sut.createQClient('eu-central-1', 'https://amazon.com/', conn)
244+
245+
assert.deepStrictEqual(client.config.region, 'eu-central-1')
246+
assert.deepStrictEqual(client.endpoint.href, 'https://amazon.com/')
247+
})
248+
})
249+
})

‎packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -384,7 +384,7 @@ describe('getChatAuthState()', function () {
384384

385385
const result = await authUtil.getChatAuthState()
386386
assert.deepStrictEqual(result, {
387-
codewhispererCore: AuthStates.connected,
387+
codewhispererCore: AuthStates.pendingProfileSelection,
388388
codewhispererChat: AuthStates.expired,
389389
amazonQ: AuthStates.expired,
390390
})
@@ -399,9 +399,9 @@ describe('getChatAuthState()', function () {
399399

400400
const result = await authUtil.getChatAuthState()
401401
assert.deepStrictEqual(result, {
402-
codewhispererCore: AuthStates.connected,
403-
codewhispererChat: AuthStates.connected,
404-
amazonQ: AuthStates.connected,
402+
codewhispererCore: AuthStates.pendingProfileSelection,
403+
codewhispererChat: AuthStates.pendingProfileSelection,
404+
amazonQ: AuthStates.pendingProfileSelection,
405405
})
406406
})
407407

‎packages/amazonq/test/unit/codewhispererChat/controllers/chat/chatRequest/converter.test.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ describe('triggerPayloadToChatRequest', () => {
4040
userInputContextLength: 0,
4141
focusFileContextLength: 0,
4242
},
43+
profile: undefined,
4344
context: [],
4445
documentReferences: [],
4546
query: undefined,

‎packages/core/package.nls.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,6 +265,7 @@
265265
"AWS.command.codewhisperer.signout": "Sign Out",
266266
"AWS.command.codewhisperer.reconnect": "Reconnect",
267267
"AWS.command.codewhisperer.openReferencePanel": "Open Code Reference Log",
268+
"AWS.command.q.selectRegionProfile": "Select Profile",
268269
"AWS.command.q.transform.acceptChanges": "Accept",
269270
"AWS.command.q.transform.rejectChanges": "Reject",
270271
"AWS.command.q.transform.stopJobInHub": "Stop job",

‎packages/core/src/amazonq/commons/baseChatStorage.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,4 +35,8 @@ export abstract class BaseChatSessionStorage<T extends { isAuthenticating: boole
3535
public deleteSession(tabID: string) {
3636
this.sessions.delete(tabID)
3737
}
38+
39+
public deleteAllSessions() {
40+
this.sessions.clear()
41+
}
3842
}

‎packages/core/src/amazonq/webview/generators/webViewContent.ts‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import { AuthUtil } from '../../../codewhisperer/util/authUtil'
99
import { FeatureConfigProvider, FeatureContext } from '../../../shared/featureConfig'
1010
import globals from '../../../shared/extensionGlobals'
1111
import { isSageMaker } from '../../../shared/extensionUtilities'
12+
import { RegionProfile } from '../../../codewhisperer/models/model'
1213

1314
export class WebViewContentGenerator {
1415
private async generateFeatureConfigsData(): Promise<string> {
@@ -85,14 +86,32 @@ export class WebViewContentGenerator {
8586
const welcomeLoadCount = globals.globalState.tryGet('aws.amazonq.welcomeChatShowCount', Number, 0)
8687
const isSMUS = isSageMaker('SMUS')
8788

89+
// only show profile card when the two conditions
90+
// 1. profile count >= 2
91+
// 2. not default (fallback) which has empty arn
92+
let regionProfile: RegionProfile | undefined = AuthUtil.instance.regionProfileManager.activeRegionProfile
93+
if (AuthUtil.instance.regionProfileManager.profiles.length === 1) {
94+
regionProfile = undefined
95+
}
96+
97+
const regionProfileString: string = JSON.stringify(regionProfile)
98+
const authState = (await AuthUtil.instance.getChatAuthState()).amazonQ
99+
88100
return `
89101
<script type="text/javascript" src="${javascriptEntrypoint.toString()}" defer onload="init()"></script>
90102
${cssLinks}
91103
<script type="text/javascript">
92104
const init = () => {
93-
createMynahUI(acquireVsCodeApi(), ${
94-
(await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected'
95-
},${featureConfigsString},${welcomeLoadCount},${disclaimerAcknowledged},${disabledCommandsString},${isSMUS});
105+
createMynahUI(
106+
acquireVsCodeApi(),
107+
${authState === 'connected'},
108+
${featureConfigsString},
109+
${welcomeLoadCount},
110+
${disclaimerAcknowledged},
111+
${regionProfileString},
112+
${disabledCommandsString},
113+
${isSMUS}
114+
);
96115
}
97116
</script>
98117
`

0 commit comments

Comments
 (0)