Skip to content

Commit c0c4764

Browse files
yuxianrzyuxianrz
committed
disable inline chat, fix other bugs, autofill access key
1 parent 3ff9966 commit c0c4764

File tree

11 files changed

+79
-29
lines changed

11 files changed

+79
-29
lines changed

‎packages/amazonq/src/extensionNode.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,8 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
100100
async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
101101
const state = AuthUtil.instance.getAuthState()
102102

103-
if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || isSageMaker())) {
104-
getLogger().error('Current Amazon Q connection is not SSO')
103+
if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || AuthUtil.instance.isIamSession() || isSageMaker())) {
104+
getLogger().error('Current Amazon Q connection is not SSO nor IAM')
105105
}
106106

107107
return {

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ export class InlineChatProvider {
143143
private async generateResponse(
144144
triggerPayload: TriggerPayload & { projectContextQueryLatencyMs?: number },
145145
triggerID: string
146-
) {
146+
): Promise<GenerateAssistantResponseCommandOutput | undefined> {
147147
const triggerEvent = this.triggerEventsStorage.getTriggerEvent(triggerID)
148148
if (triggerEvent === undefined) {
149149
return
@@ -182,7 +182,12 @@ export class InlineChatProvider {
182182
let response: GenerateAssistantResponseCommandOutput | undefined = undefined
183183
session.createNewTokenSource()
184184
try {
185-
response = await session.chatSso(request)
185+
if (AuthUtil.instance.isSsoSession()) {
186+
response = await session.chatSso(request)
187+
} else {
188+
// Call sendMessage because Q Developer Streaming Client does not have generateAssistantResponse
189+
throw new ToolkitError('Inline chat is only available with SSO authentication')
190+
}
186191
getLogger().info(
187192
`response to tab: ${tabID} conversationID: ${session.sessionIdentifier} requestID: ${response.$metadata.requestId} metadata: %O`,
188193
response.$metadata

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -483,7 +483,7 @@ describe('AuthUtil', async function () {
483483
await auth.getIamCredential()
484484
assert.fail('Should have thrown an error')
485485
} catch (err) {
486-
assert.strictEqual((err as Error).message, 'Cannot get token with SSO session')
486+
assert.strictEqual((err as Error).message, 'Cannot get credential with SSO session')
487487
}
488488
})
489489

‎packages/core/src/codewhisperer/ui/codeWhispererNodes.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,11 @@ export function createManageSubscription(): DataQuickPickItem<'manageSubscriptio
192192
export function createSignout(): DataQuickPickItem<'signout'> {
193193
const label = localize('AWS.codewhisperer.signoutNode.label', 'Sign Out')
194194
const icon = getIcon('vscode-export')
195-
const connection = AuthUtil.instance.isBuilderIdConnection() ? 'AWS Builder ID' : 'IAM Identity Center'
195+
const connection = AuthUtil.instance.isIamConnection()
196+
? 'IAM Credentials'
197+
: AuthUtil.instance.isBuilderIdConnection()
198+
? 'AWS Builder ID'
199+
: 'IAM Identity Center'
196200

197201
return {
198202
data: 'signout',

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -219,7 +219,7 @@ export class AuthUtil implements IAuthProvider {
219219
if (this.session) {
220220
const credential = (await this.session.getCredential()).credential
221221
if (typeof credential !== 'object') {
222-
throw new ToolkitError('Cannot get token with SSO session')
222+
throw new ToolkitError('Cannot get credential with SSO session')
223223
}
224224
return credential
225225
} else {

‎packages/core/src/login/webview/vue/backend.ts‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import { getLogger } from '../../../shared/logger/logger'
3333
import { isValidUrl } from '../../../shared/utilities/uriUtils'
3434
import { RegionProfile } from '../../../codewhisperer/models/model'
3535
import { ProfileSwitchIntent } from '../../../codewhisperer/region/regionProfileManager'
36+
import { showMessage } from '../../../shared/utilities/messages'
3637

3738
export abstract class CommonAuthWebview extends VueWebview {
3839
private readonly className = 'CommonAuthWebview'
@@ -183,7 +184,7 @@ export abstract class CommonAuthWebview extends VueWebview {
183184
abstract fetchConnections(): Promise<AwsConnection[] | undefined>
184185

185186
async errorNotification(e: AuthError) {
186-
void vscode.window.showInformationMessage(`${e.text}`)
187+
showMessage('error', e.text)
187188
}
188189

189190
abstract quitLoginScreen(): Promise<void>
@@ -296,6 +297,15 @@ export abstract class CommonAuthWebview extends VueWebview {
296297
return globals.globalState.tryGet('recentSso', Object, { startUrl: '', region: 'us-east-1' })
297298
}
298299

300+
getDefaultIamKeys(): { accessKey: string } {
301+
const devSettings = DevSettings.instance.get('autofillAccessKey', '')
302+
if (devSettings) {
303+
return { accessKey: devSettings }
304+
}
305+
306+
return globals.globalState.tryGet('recentIamKeys', Object, { accessKey: '' })
307+
}
308+
299309
cancelAuthFlow() {
300310
AuthSSOServer.lastInstance?.cancelCurrentFlow()
301311
}

‎packages/core/src/login/webview/vue/login.vue‎

Lines changed: 42 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -230,7 +230,7 @@
230230

231231
<template v-if="stage === 'AUTHENTICATING'">
232232
<div class="auth-container-section">
233-
<div v-if="app === 'TOOLKIT' && profileName.length > 0" class="header bottomMargin">
233+
<div v-if="selectedLoginOption === LoginOption.IAM_CREDENTIAL" class="header bottomMargin">
234234
Connecting to IAM...
235235
</div>
236236
<div v-else class="header bottomMargin">Authenticating in browser...</div>
@@ -273,7 +273,7 @@
273273
<div class="title">Secret Access Key</div>
274274
<input
275275
class="iamInput bottomMargin"
276-
type="text"
276+
type="password"
277277
id="secretKey"
278278
name="secretKey"
279279
v-model="secretKey"
@@ -330,6 +330,10 @@ interface ImportedLogin {
330330
type: number
331331
startUrl: string
332332
region: string
333+
// Add IAM credential fields
334+
profileName?: string
335+
accessKey?: string
336+
secretKey?: string // Note: storing secrets has security implications
333337
}
334338
335339
export default defineComponent({
@@ -349,6 +353,7 @@ export default defineComponent({
349353
data() {
350354
return {
351355
existingStartUrls: [] as string[],
356+
existingIamAccessKeys: [] as string[],
352357
importedLogins: [] as ImportedLogin[],
353358
selectedLoginOption: LoginOption.NONE,
354359
stage: 'START' as Stage,
@@ -368,6 +373,8 @@ export default defineComponent({
368373
const defaultSso = await this.getDefaultSso()
369374
this.startUrl = defaultSso.startUrl
370375
this.selectedRegion = defaultSso.region
376+
const defaultIamAccessKey = await this.getDefaultIamAccessKey()
377+
this.accessKey = defaultIamAccessKey.accessKey
371378
await this.emitUpdate('created')
372379
},
373380
@@ -397,6 +404,10 @@ export default defineComponent({
397404
}
398405
},
399406
handleDocumentClick(event: any) {
407+
// Only reset selection when in START stage to avoid clearing during authentication
408+
if (this.stage !== 'START') {
409+
return
410+
}
400411
const isClickInsideSelectableItems = event.target.closest('.selectable-item')
401412
if (!isClickInsideSelectableItems) {
402413
this.selectedLoginOption = 0
@@ -437,17 +448,32 @@ export default defineComponent({
437448
const selectedConnection =
438449
this.importedLogins[this.selectedLoginOption - LoginOption.IMPORTED_LOGINS]
439450
440-
// Imported connections cannot be Builder IDs, they are filtered out in the client.
441-
const error = await client.startEnterpriseSetup(
442-
selectedConnection.startUrl,
443-
selectedConnection.region,
444-
this.app
445-
)
446-
if (error) {
447-
this.stage = 'START'
448-
void client.errorNotification(error)
449-
} else {
450-
this.stage = 'CONNECTED'
451+
// Handle both SSO and IAM imported connections
452+
if (selectedConnection.type === LoginOption.ENTERPRISE_SSO) {
453+
const error = await client.startEnterpriseSetup(
454+
selectedConnection.startUrl,
455+
selectedConnection.region,
456+
this.app
457+
)
458+
if (error) {
459+
this.stage = 'START'
460+
void client.errorNotification(error)
461+
} else {
462+
this.stage = 'CONNECTED'
463+
}
464+
} else if (selectedConnection.type === LoginOption.IAM_CREDENTIAL) {
465+
// Use stored IAM credentials
466+
const error = await client.startIamCredentialSetup(
467+
selectedConnection.profileName || '',
468+
selectedConnection.accessKey || '',
469+
selectedConnection.secretKey || ''
470+
)
471+
if (error) {
472+
this.stage = 'START'
473+
void client.errorNotification(error)
474+
} else {
475+
this.stage = 'CONNECTED'
476+
}
451477
}
452478
} else if (this.selectedLoginOption === LoginOption.IAM_CREDENTIAL) {
453479
this.stage = 'AWS_PROFILE'
@@ -581,6 +607,9 @@ export default defineComponent({
581607
async getDefaultSso() {
582608
return await client.getDefaultSsoProfile()
583609
},
610+
async getDefaultIamAccessKey() {
611+
return await client.getDefaultIamKeys()
612+
},
584613
handleHelpLinkClick() {
585614
void client.emitUiClick('auth_helpLink')
586615
},

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ export class FeatureConfigProvider {
118118
}
119119

120120
async fetchFeatureConfigs(): Promise<void> {
121-
if (AuthUtil.instance.isConnectionExpired()) {
121+
if (AuthUtil.instance.isConnectionExpired() || AuthUtil.instance.isIamSession()) {
122122
return
123123
}
124124

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ export type globalKey =
7171
| 'lastOsStartTime'
7272
| 'recentCredentials'
7373
| 'recentSso'
74+
| 'recentIamKeys'
7475
// List of regions enabled in AWS Explorer.
7576
| 'region'
7677
// TODO: implement this via `PromptSettings` instead of globalState.

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -780,6 +780,7 @@ const devSettings = {
780780
amazonqWorkspaceLsp: Record(String, String),
781781
ssoCacheDirectory: String,
782782
autofillStartUrl: String,
783+
autofillAccessKey: String,
783784
webAuth: Boolean,
784785
notificationsPollInterval: Number,
785786
}

0 commit comments

Comments
 (0)