Skip to content

Commit 46b2bdf

Browse files
opieter-awsopieter-aws
committed
feat(core): Update all references to use new auth class using flare identity server for Amazon Q
1 parent 528fa0d commit 46b2bdf

File tree

76 files changed

+901
-999
lines changed

Some content is hidden

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

76 files changed

+901
-999
lines changed

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import { Messenger } from './chat/controller/messenger/messenger'
1919
import { UIMessageListener } from './chat/views/actions/uiMessageListener'
2020
import { debounce } from 'lodash'
2121
import { Commands, placeholder } from 'aws-core-vscode/shared'
22+
import { auth2 } from 'aws-core-vscode/auth'
2223

2324
export function init(appContext: AmazonQAppInitContext) {
2425
const scanChatControllerEventEmitters: ScanChatControllerEventEmitters = {
@@ -52,7 +53,7 @@ export function init(appContext: AmazonQAppInitContext) {
5253
appContext.registerWebViewToAppMessagePublisher(new MessagePublisher<any>(scanChatUIInputEventEmitter), 'review')
5354

5455
const debouncedEvent = debounce(async () => {
55-
const authenticated = (await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected'
56+
const authenticated = AuthUtil.instance.getAuthState() === 'connected'
5657
let authenticatingSessionID = ''
5758

5859
if (authenticated) {
@@ -67,7 +68,7 @@ export function init(appContext: AmazonQAppInitContext) {
6768
messenger.sendAuthenticationUpdate(authenticated, [authenticatingSessionID])
6869
}, 500)
6970

70-
AuthUtil.instance.secondaryAuth.onDidChangeActiveConnection(() => {
71+
AuthUtil.instance.onDidChangeConnectionState((e: auth2.AuthStateEvent) => {
7172
return debouncedEvent()
7273
})
7374
AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(() => {

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ export class ScanController {
104104
telemetry.amazonq_feedback.emit({
105105
featureId: 'amazonQReview',
106106
amazonqConversationId: this.sessionStorage.getSession().scanUuid,
107-
credentialStartUrl: AuthUtil.instance.startUrl,
107+
credentialStartUrl: AuthUtil.instance.connection?.startUrl,
108108
interactionType: data.vote,
109109
})
110110
})
@@ -122,8 +122,8 @@ export class ScanController {
122122
try {
123123
getLogger().debug(`Q - Review: Session created with id: ${session.tabID}`)
124124

125-
const authState = await AuthUtil.instance.getChatAuthState()
126-
if (authState.amazonQ !== 'connected') {
125+
const authState = AuthUtil.instance.getAuthState()
126+
if (authState !== 'connected') {
127127
void this.messenger.sendAuthNeededExceptionMessage(authState, tabID)
128128
session.isAuthenticating = true
129129
return
@@ -161,8 +161,8 @@ export class ScanController {
161161
return
162162
}
163163
// check that the session is authenticated
164-
const authState = await AuthUtil.instance.getChatAuthState()
165-
if (authState.amazonQ !== 'connected') {
164+
const authState = AuthUtil.instance.getAuthState()
165+
if (authState !== 'connected') {
166166
void this.messenger.sendAuthNeededExceptionMessage(authState, message.tabID)
167167
session.isAuthenticating = true
168168
return

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

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010

1111
import { AuthFollowUpType, AuthMessageDataMap } from 'aws-core-vscode/amazonq'
1212
import {
13-
FeatureAuthState,
1413
SecurityScanError,
1514
CodeWhispererConstants,
1615
SecurityScanStep,
@@ -34,6 +33,7 @@ import {
3433
import { i18n } from 'aws-core-vscode/shared'
3534
import { ScanAction, scanProgressMessage } from '../../../models/constants'
3635
import path from 'path'
36+
import { auth2 } from 'aws-core-vscode/auth'
3737

3838
export type UnrecoverableErrorType = 'no-project-found' | 'no-open-file-found' | 'invalid-file-type'
3939

@@ -78,19 +78,15 @@ export class Messenger {
7878
this.dispatcher.sendUpdatePromptProgress(new UpdatePromptProgressMessage(tabID, progressField))
7979
}
8080

81-
public async sendAuthNeededExceptionMessage(credentialState: FeatureAuthState, tabID: string) {
81+
public async sendAuthNeededExceptionMessage(credentialState: auth2.AuthState, tabID: string) {
8282
let authType: AuthFollowUpType = 'full-auth'
8383
let message = AuthMessageDataMap[authType].message
8484

85-
switch (credentialState.amazonQ) {
86-
case 'disconnected':
85+
switch (credentialState) {
86+
case 'notConnected':
8787
authType = 'full-auth'
8888
message = AuthMessageDataMap[authType].message
8989
break
90-
case 'unsupported':
91-
authType = 'use-supported-auth'
92-
message = AuthMessageDataMap[authType].message
93-
break
9490
case 'expired':
9591
authType = 're-auth'
9692
message = AuthMessageDataMap[authType].message

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

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -123,10 +123,8 @@ export class InlineChatProvider {
123123

124124
const tabID = triggerEvent.tabID
125125

126-
const credentialsState = await AuthUtil.instance.getChatAuthState()
127-
if (
128-
!(credentialsState.codewhispererChat === 'connected' && credentialsState.codewhispererCore === 'connected')
129-
) {
126+
const credentialsState = AuthUtil.instance.getAuthState()
127+
if (credentialsState !== 'connected') {
130128
const { message } = extractAuthFollowUp(credentialsState)
131129
this.errorEmitter.fire()
132130
throw new ToolkitError(message)

‎packages/amazonq/src/lsp/chat/messages.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ export function registerMessageListeners(
117117

118118
if (fullAuthTypes.includes(authType)) {
119119
try {
120-
await AuthUtil.instance.secondaryAuth.deleteConnection()
120+
await AuthUtil.instance.logout()
121121
} catch (e) {
122122
languageClient.error(
123123
`[VSCode Client] Failed to authenticate after AUTH_FOLLOW_UP_CLICKED: ${(e as Error).message}`

‎packages/amazonq/test/e2e/amazonq/utils/setup.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,9 @@
55
import { AuthUtil } from 'aws-core-vscode/codewhisperer'
66

77
export async function loginToIdC() {
8-
const authState = await AuthUtil.instance.getChatAuthState()
8+
const authState = AuthUtil.instance.getAuthState()
99
if (process.env['AWS_TOOLKIT_AUTOMATION'] === 'local') {
10-
if (authState.amazonQ !== 'connected') {
10+
if (authState !== 'connected') {
1111
throw new Error('You will need to login manually before running tests.')
1212
}
1313
return
@@ -22,5 +22,5 @@ export async function loginToIdC() {
2222
)
2323
}
2424

25-
await AuthUtil.instance.connectToEnterpriseSso(startUrl, region)
25+
await AuthUtil.instance.login(startUrl, region)
2626
}

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

Lines changed: 49 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -7,41 +7,36 @@ import * as sinon from 'sinon'
77
import assert, { fail } from 'assert'
88
import { AuthUtil, RegionProfile, RegionProfileManager, defaultServiceConfig } from 'aws-core-vscode/codewhisperer'
99
import { globals } from 'aws-core-vscode/shared'
10-
import { createTestAuth } from 'aws-core-vscode/test'
11-
import { SsoConnection } from 'aws-core-vscode/auth'
10+
import { builderIdStartUrl } from 'aws-core-vscode/auth'
1211

1312
const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start'
13+
const region = 'us-east-1'
1414

1515
describe('RegionProfileManager', function () {
16-
let sut: RegionProfileManager
17-
let auth: ReturnType<typeof createTestAuth>
18-
let authUtil: AuthUtil
16+
let regionProfileManager: RegionProfileManager
1917

2018
const profileFoo: RegionProfile = {
2119
name: 'foo',
22-
region: 'us-east-1',
20+
region,
2321
arn: 'foo arn',
2422
description: 'foo description',
2523
}
2624

2725
async function setupConnection(type: 'builderId' | 'idc') {
2826
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')
27+
await AuthUtil.instance.login(builderIdStartUrl, region)
28+
assert.ok(AuthUtil.instance.isSsoSession())
29+
assert.ok(AuthUtil.instance.isBuilderIdConnection())
3330
} 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)')
31+
await AuthUtil.instance.login(enterpriseSsoStartUrl, region)
32+
assert.ok(AuthUtil.instance.isSsoSession())
33+
assert.ok(AuthUtil.instance.isIdcConnection())
3834
}
3935
}
4036

4137
beforeEach(function () {
42-
auth = createTestAuth(globals.globalState)
43-
authUtil = new AuthUtil(auth)
44-
sut = new RegionProfileManager(() => authUtil.conn)
38+
regionProfileManager = new RegionProfileManager(AuthUtil.instance)
39+
// const authUtilStub = sinon.stub(AuthUtil.instance, 'isIdcConnection').returns(isSso)
4540
})
4641

4742
afterEach(function () {
@@ -65,9 +60,9 @@ describe('RegionProfileManager', function () {
6560
const mockClient = {
6661
listAvailableProfiles: listProfilesStub,
6762
}
68-
const createClientStub = sinon.stub(sut, 'createQClient').resolves(mockClient)
63+
const createClientStub = sinon.stub(regionProfileManager, 'createQClient').resolves(mockClient)
6964

70-
const r = await sut.listRegionProfile()
65+
const r = await regionProfileManager.listRegionProfiles()
7166

7267
assert.strictEqual(r.length, 2)
7368
assert.deepStrictEqual(r, [
@@ -93,41 +88,40 @@ describe('RegionProfileManager', function () {
9388
describe('switch and get profile', function () {
9489
it('should switch if connection is IdC', async function () {
9590
await setupConnection('idc')
96-
await sut.switchRegionProfile(profileFoo, 'user')
97-
assert.deepStrictEqual(sut.activeRegionProfile, profileFoo)
91+
await regionProfileManager.switchRegionProfile(profileFoo, 'user')
92+
assert.deepStrictEqual(regionProfileManager.activeRegionProfile, profileFoo)
9893
})
9994

10095
it('should do nothing and return undefined if connection is builder id', async function () {
10196
await setupConnection('builderId')
102-
await sut.switchRegionProfile(profileFoo, 'user')
103-
assert.deepStrictEqual(sut.activeRegionProfile, undefined)
97+
await regionProfileManager.switchRegionProfile(profileFoo, 'user')
98+
assert.deepStrictEqual(regionProfileManager.activeRegionProfile, undefined)
10499
})
105100
})
106101

107102
describe(`client config`, function () {
108103
it(`no valid credential should throw`, async function () {
109-
assert.ok(authUtil.conn === undefined)
104+
assert.ok(!AuthUtil.instance.isConnected())
110105

111106
assert.throws(() => {
112-
sut.clientConfig
107+
regionProfileManager.clientConfig
113108
}, /trying to get client configuration without credential/)
114109
})
115110

116111
it(`builder id should always use default profile IAD`, async function () {
117112
await setupConnection('builderId')
118-
await sut.switchRegionProfile(profileFoo, 'user')
119-
assert.deepStrictEqual(sut.activeRegionProfile, undefined)
120-
const conn = authUtil.conn
121-
if (!conn) {
113+
await regionProfileManager.switchRegionProfile(profileFoo, 'user')
114+
assert.deepStrictEqual(regionProfileManager.activeRegionProfile, undefined)
115+
if (!AuthUtil.instance.isConnected()) {
122116
fail('connection should not be undefined')
123117
}
124118

125-
assert.deepStrictEqual(sut.clientConfig, defaultServiceConfig)
119+
assert.deepStrictEqual(regionProfileManager.clientConfig, defaultServiceConfig)
126120
})
127121

128122
it(`idc should return correct endpoint corresponding to profile region`, async function () {
129123
await setupConnection('idc')
130-
await sut.switchRegionProfile(
124+
await regionProfileManager.switchRegionProfile(
131125
{
132126
name: 'foo',
133127
region: 'eu-central-1',
@@ -136,16 +130,16 @@ describe('RegionProfileManager', function () {
136130
},
137131
'user'
138132
)
139-
assert.ok(sut.activeRegionProfile)
140-
assert.deepStrictEqual(sut.clientConfig, {
133+
assert.ok(regionProfileManager.activeRegionProfile)
134+
assert.deepStrictEqual(regionProfileManager.clientConfig, {
141135
region: 'eu-central-1',
142136
endpoint: 'https://q.eu-central-1.amazonaws.com/',
143137
})
144138
})
145139

146140
it(`idc should throw if corresponding endpoint is not defined`, async function () {
147141
await setupConnection('idc')
148-
await sut.switchRegionProfile(
142+
await regionProfileManager.switchRegionProfile(
149143
{
150144
name: 'foo',
151145
region: 'unknown region',
@@ -156,74 +150,71 @@ describe('RegionProfileManager', function () {
156150
)
157151

158152
assert.throws(() => {
159-
sut.clientConfig
153+
regionProfileManager.clientConfig
160154
}, /Q client configuration error, endpoint not found for region*/)
161155
})
162156
})
163157

164158
describe('persistence', function () {
165159
it('persistSelectedRegionProfile', async function () {
166160
await setupConnection('idc')
167-
await sut.switchRegionProfile(profileFoo, 'user')
168-
assert.deepStrictEqual(sut.activeRegionProfile, profileFoo)
169-
const conn = authUtil.conn
170-
if (!conn) {
161+
await regionProfileManager.switchRegionProfile(profileFoo, 'user')
162+
assert.deepStrictEqual(regionProfileManager.activeRegionProfile, profileFoo)
163+
if (!AuthUtil.instance.isConnected()) {
171164
fail('connection should not be undefined')
172165
}
173166

174-
await sut.persistSelectRegionProfile()
167+
await regionProfileManager.persistSelectRegionProfile()
175168

176169
const state = globals.globalState.tryGet<{ [label: string]: RegionProfile }>(
177170
'aws.amazonq.regionProfiles',
178171
Object,
179172
{}
180173
)
181174

182-
assert.strictEqual(state[conn.id], profileFoo)
175+
assert.strictEqual(state[AuthUtil.instance.profileName], profileFoo)
183176
})
184177

185178
it(`restoreRegionProfile`, async function () {
186-
sinon.stub(sut, 'listRegionProfile').resolves([profileFoo])
179+
sinon.stub(regionProfileManager, 'listRegionProfiles').resolves([profileFoo])
187180
await setupConnection('idc')
188-
const conn = authUtil.conn
189-
if (!conn) {
181+
if (!AuthUtil.instance.isConnected()) {
190182
fail('connection should not be undefined')
191183
}
192184

193185
const state = {} as any
194-
state[conn.id] = profileFoo
186+
state[AuthUtil.instance.profileName] = profileFoo
195187

196188
await globals.globalState.update('aws.amazonq.regionProfiles', state)
197189

198-
await sut.restoreRegionProfile(conn)
190+
await regionProfileManager.restoreRegionProfile()
199191

200-
assert.strictEqual(sut.activeRegionProfile, profileFoo)
192+
assert.strictEqual(regionProfileManager.activeRegionProfile, profileFoo)
201193
})
202194
})
203195

204196
describe('invalidate', function () {
205197
it('should reset activeProfile and global state', async function () {
206198
// setup
207199
await setupConnection('idc')
208-
await sut.switchRegionProfile(profileFoo, 'user')
209-
assert.deepStrictEqual(sut.activeRegionProfile, profileFoo)
210-
const conn = authUtil.conn
211-
if (!conn) {
200+
await regionProfileManager.switchRegionProfile(profileFoo, 'user')
201+
assert.deepStrictEqual(regionProfileManager.activeRegionProfile, profileFoo)
202+
if (!AuthUtil.instance.isConnected()) {
212203
fail('connection should not be undefined')
213204
}
214-
await sut.persistSelectRegionProfile()
205+
await regionProfileManager.persistSelectRegionProfile()
215206
const state = globals.globalState.tryGet<{ [label: string]: RegionProfile }>(
216207
'aws.amazonq.regionProfiles',
217208
Object,
218209
{}
219210
)
220-
assert.strictEqual(state[conn.id], profileFoo)
211+
assert.strictEqual(state[AuthUtil.instance.profileName], profileFoo)
221212

222213
// subject to test
223-
await sut.invalidateProfile(profileFoo.arn)
214+
await regionProfileManager.invalidateProfile(profileFoo.arn)
224215

225216
// assertion
226-
assert.strictEqual(sut.activeRegionProfile, undefined)
217+
assert.strictEqual(regionProfileManager.activeRegionProfile, undefined)
227218
const actualGlobalState = globals.globalState.tryGet<{ [label: string]: RegionProfile }>(
228219
'aws.amazonq.regionProfiles',
229220
Object,
@@ -236,11 +227,10 @@ describe('RegionProfileManager', function () {
236227
describe('createQClient', function () {
237228
it(`should configure the endpoint and region correspondingly`, async function () {
238229
await setupConnection('idc')
239-
await sut.switchRegionProfile(profileFoo, 'user')
240-
assert.deepStrictEqual(sut.activeRegionProfile, profileFoo)
241-
const conn = authUtil.conn as SsoConnection
230+
await regionProfileManager.switchRegionProfile(profileFoo, 'user')
231+
assert.deepStrictEqual(regionProfileManager.activeRegionProfile, profileFoo)
242232

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

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

‎packages/amazonq/test/unit/codewhisperer/service/codewhisperer.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,7 @@ describe('codewhisperer', async function () {
136136
}),
137137
} as Request<SendTelemetryEventResponse, AWSError>)
138138

139-
const authUtilStub = sinon.stub(AuthUtil.instance, 'isValidEnterpriseSsoInUse').returns(isSso)
139+
const authUtilStub = sinon.stub(AuthUtil.instance, 'isIdcConnection').returns(isSso)
140140
await globals.telemetry.setTelemetryEnabled(isTelemetryEnabled)
141141
await codeWhispererClient.sendTelemetryEvent({ telemetryEvent: payload })
142142
const expectedOptOutPreference = isTelemetryEnabled ? 'OPTIN' : 'OPTOUT'

0 commit comments

Comments
 (0)