Skip to content

Commit 528fa0d

Browse files
opieter-awsopieter-aws
committed
Merge remote-tracking branch 'upstream/feature/amazonqLSP-auth' into feature/amazonqLSP-auth
2 parents ea36f11 + b896e4f commit 528fa0d

File tree

3 files changed

+60
-36
lines changed

3 files changed

+60
-36
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ export class AmazonQChatViewProvider implements WebviewViewProvider {
6464
}
6565

6666
private async getWebviewContent(mynahUIPath: string) {
67-
const disclaimerAcknowledged = AmazonQPromptSettings.instance.isPromptEnabled('amazonQChatDisclaimer')
67+
const disclaimerAcknowledged = !AmazonQPromptSettings.instance.isPromptEnabled('amazonQChatDisclaimer')
6868
return `
6969
<!DOCTYPE html>
7070
<html lang="en">

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

Lines changed: 11 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -94,20 +94,11 @@ export const baseCustomization = {
9494
* @returns customization selected by users, `baseCustomization` if none is selected
9595
*/
9696
export const getSelectedCustomization = (): Customization => {
97-
if (
98-
!AuthUtil.instance.isCustomizationFeatureEnabled ||
99-
!AuthUtil.instance.isValidEnterpriseSsoInUse() ||
100-
!AuthUtil.instance.conn
101-
) {
97+
if (!AuthUtil.instance.isCustomizationFeatureEnabled || !AuthUtil.instance.isIdcConnection()) {
10298
return baseCustomization
10399
}
104100

105-
const selectedCustomizationArr = globals.globalState.tryGet<{ [label: string]: Customization }>(
106-
'CODEWHISPERER_SELECTED_CUSTOMIZATION',
107-
Object,
108-
{}
109-
)
110-
const selectedCustomization = selectedCustomizationArr[AuthUtil.instance.conn.label]
101+
const selectedCustomization = globals.globalState.getAmazonQCustomization(AuthUtil.instance.profileName)
111102

112103
if (selectedCustomization && selectedCustomization.name !== '') {
113104
return selectedCustomization
@@ -124,7 +115,7 @@ export const getSelectedCustomization = (): Customization => {
124115
* 2. the override customization arn is different from the previous override customization if any. The purpose is to only do override once on users' behalf.
125116
*/
126117
export const setSelectedCustomization = async (customization: Customization, isOverride: boolean = false) => {
127-
if (!AuthUtil.instance.isValidEnterpriseSsoInUse() || !AuthUtil.instance.conn) {
118+
if (!AuthUtil.instance.isIdcConnection()) {
128119
return
129120
}
130121
if (isOverride) {
@@ -133,15 +124,10 @@ export const setSelectedCustomization = async (customization: Customization, isO
133124
return
134125
}
135126
}
136-
const selectedCustomizationObj = globals.globalState.tryGet<{ [label: string]: Customization }>(
137-
'CODEWHISPERER_SELECTED_CUSTOMIZATION',
138-
Object,
139-
{}
140-
)
141-
selectedCustomizationObj[AuthUtil.instance.conn.label] = customization
142-
getLogger().debug(`Selected customization ${customization.name} for ${AuthUtil.instance.conn.label}`)
143127

144-
await globals.globalState.update('CODEWHISPERER_SELECTED_CUSTOMIZATION', selectedCustomizationObj)
128+
await globals.globalState.update('CODEWHISPERER_SELECTED_CUSTOMIZATION', customization)
129+
getLogger().debug(`Selected customization ${customization.name} for ${AuthUtil.instance.profileName}`)
130+
145131
if (isOverride) {
146132
await globals.globalState.update('aws.amazonq.customization.overrideV2', customization.arn)
147133
}
@@ -150,28 +136,18 @@ export const setSelectedCustomization = async (customization: Customization, isO
150136
}
151137

152138
export const getPersistedCustomizations = (): Customization[] => {
153-
if (!AuthUtil.instance.isValidEnterpriseSsoInUse() || !AuthUtil.instance.conn) {
139+
if (!AuthUtil.instance.isIdcConnection()) {
154140
return []
155141
}
156-
const persistedCustomizationsObj = globals.globalState.tryGet<{ [label: string]: Customization[] }>(
157-
'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS',
158-
Object,
159-
{}
160-
)
161-
return persistedCustomizationsObj[AuthUtil.instance.conn.label] || []
142+
return globals.globalState.getAmazonQCachedCustomization(AuthUtil.instance.profileName)
162143
}
163144

164145
export const setPersistedCustomizations = async (customizations: Customization[]) => {
165-
if (!AuthUtil.instance.isValidEnterpriseSsoInUse() || !AuthUtil.instance.conn) {
146+
if (!AuthUtil.instance.isIdcConnection()) {
166147
return
167148
}
168-
const persistedCustomizationsObj = globals.globalState.tryGet<{ [label: string]: Customization[] }>(
169-
'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS',
170-
Object,
171-
{}
172-
)
173-
persistedCustomizationsObj[AuthUtil.instance.conn.label] = customizations
174-
await globals.globalState.update('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', persistedCustomizationsObj)
149+
150+
await globals.globalState.update('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', customizations)
175151
}
176152

177153
export const getNewCustomizationsAvailable = () => {

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

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import * as vscode from 'vscode'
77
import { getLogger } from './logger/logger'
88
import * as redshift from '../awsService/redshift/models/models'
99
import { TypeConstructor, cast } from './utilities/typeConstructors'
10+
import { Customization } from '../codewhisperer/client/codewhispereruserclient'
1011

1112
type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId'
1213
export type ToolIdStateKey = `${ToolId}.savedConnectionId`
@@ -227,6 +228,53 @@ export class GlobalState implements vscode.Memento {
227228
return all?.[warehouseArn]
228229
}
229230

231+
/**
232+
* Get the Amazon Q customization. If legacy (map of customizations) store the
233+
* customization with label of profile name
234+
*
235+
* @param profileName name of profile, only used in case legacy customization is found
236+
* @returns Amazon Q customization, or undefined if not found.
237+
* If legacy, return the Amazon Q customization for the auth profile name
238+
*/
239+
getAmazonQCustomization(profileName: string): Customization | undefined {
240+
const result = this.tryGet('CODEWHISPERER_SELECTED_CUSTOMIZATION', Object, undefined)
241+
242+
// Legacy migration for old customization map of type { [label: string]: Customization[] }
243+
if (typeof result === 'object' && Object.values(result).every(Array.isArray)) {
244+
const selectedCustomization = result[profileName]
245+
this.tryUpdate('CODEWHISPERER_SELECTED_CUSTOMIZATION', selectedCustomization)
246+
return selectedCustomization
247+
} else {
248+
return result
249+
}
250+
}
251+
252+
/**
253+
* Get the Amazon Q cached customizations. If legacy (map of customizations) store the
254+
* customizations with label of profile name
255+
*
256+
* @param profileName name of profile, only used in case legacy customization is found
257+
* @returns array of Amazon Q cached customizations, or empty array if not found.
258+
* If legacy, return the Amazon Q persisted customizations for the auth profile name
259+
*/
260+
getAmazonQCachedCustomization(profileName: string): Customization[] {
261+
const result = this.tryGet<Customization[]>('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', Array, [])
262+
263+
// Legacy migration for old customization map of type { [label: string]: Customization[] }
264+
if (result.length === 0) {
265+
const customizations = this.tryGet<{ [label: string]: Customization[] }>(
266+
'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS',
267+
Object,
268+
{}
269+
)
270+
const cachedCustomizationsArray = customizations[profileName] || []
271+
this.tryUpdate('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', cachedCustomizationsArray)
272+
return cachedCustomizationsArray
273+
} else {
274+
return result
275+
}
276+
}
277+
230278
/**
231279
* Sets SSO session creation timestamp for the given session `id`.
232280
*

0 commit comments

Comments
 (0)