Skip to content

Commit 03441f1

Browse files
aws-toolkit-automationaws-toolkit-automation
authored
Merge main into feature/q-lsp
2 parents c895f6b + 951d519 commit 03441f1

File tree

4 files changed

+147
-10
lines changed

4 files changed

+147
-10
lines changed

‎plugins/amazonq/codewhisperer/jetbrains-community/src/migration/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ interface CodeWhispererModelConfigurator {
1818

1919
fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?)
2020

21+
fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?, isOverride: Boolean)
22+
2123
/**
2224
* This method is only used for invalidate a stale customization which was previously active but was removed, it will remove all usage of this customization
2325
* but not limited to the specific connection.

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,8 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
7676

7777
private val hasShownNewCustomizationNotification = AtomicBoolean(false)
7878

79+
private var serviceDefaultArn: String? = null
80+
7981
override fun showConfigDialog(project: Project) {
8082
runInEdt {
8183
calculateIfIamIdentityCenterConnection(project) {
@@ -165,20 +167,23 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
165167
override fun activeCustomization(project: Project): CodeWhispererCustomization? {
166168
val selectedCustomization = calculateIfIamIdentityCenterConnection(project) { connectionIdToActiveCustomizationArn[it.id] }
167169

168-
if (selectedCustomization != null) {
169-
return selectedCustomization
170-
} else {
171-
val customizationOverride = CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()
172-
if (customizationOverride == null || customizationOverride.value.stringValue().isEmpty()) return null
173-
return CodeWhispererCustomization(
174-
arn = customizationOverride.value.stringValue(),
175-
name = customizationOverride.variation,
176-
)
177-
}
170+
return selectedCustomization
178171
}
179172

180173
override fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?) {
174+
switchCustomization(project, newCustomization, false)
175+
}
176+
177+
/**
178+
* Override happens when ALL following conditions are met
179+
* 1. service returns non-empty override customization arn, refer to [CodeWhispererFeatureConfigService]
180+
* 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.
181+
*/
182+
override fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?, isOverride: Boolean) {
181183
calculateIfIamIdentityCenterConnection(project) {
184+
if (isOverride && (newCustomization == null || newCustomization.arn.isEmpty() || serviceDefaultArn == newCustomization.arn)) {
185+
return@calculateIfIamIdentityCenterConnection
186+
}
182187
val oldCus = connectionIdToActiveCustomizationArn[it.id]
183188
if (oldCus != newCustomization) {
184189
newCustomization?.let { newCus ->
@@ -191,6 +196,9 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
191196

192197
CodeWhispererCustomizationListener.notifyCustomUiUpdate()
193198
}
199+
if (isOverride) {
200+
serviceDefaultArn = newCustomization?.arn
201+
}
194202
}
195203
}
196204

@@ -240,6 +248,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
240248
val state = CodeWhispererCustomizationState()
241249
state.connectionIdToActiveCustomizationArn.putAll(this.connectionIdToActiveCustomizationArn)
242250
state.previousAvailableCustomizations.putAll(this.connectionToCustomizationsShownLastTime)
251+
state.serviceDefaultArn = this.serviceDefaultArn
243252

244253
return state
245254
}
@@ -250,6 +259,8 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
250259

251260
connectionToCustomizationsShownLastTime.clear()
252261
connectionToCustomizationsShownLastTime.putAll(state.previousAvailableCustomizations)
262+
263+
this.serviceDefaultArn = state.serviceDefaultArn
253264
}
254265

255266
override fun dispose() {}
@@ -280,6 +291,9 @@ class CodeWhispererCustomizationState : BaseState() {
280291
@get:Property
281292
@get:MapAnnotation
282293
val previousAvailableCustomizations by map<String, MutableList<String>>()
294+
295+
@get:Property
296+
var serviceDefaultArn by string()
283297
}
284298

285299
data class CustomizationUiItem(

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
1313
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
1414
import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
1515
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
16+
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
1617
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
1718
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
1819
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isUserBuilderId
@@ -78,6 +79,14 @@ class CodeWhispererProjectStartupActivity : StartupActivity.DumbAware {
7879
projectCoroutineScope(project).launch {
7980
while (isActive) {
8081
CodeWhispererFeatureConfigService.getInstance().fetchFeatureConfigs(project)
82+
CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.let { customization ->
83+
CodeWhispererModelConfigurator.getInstance().switchCustomization(
84+
project,
85+
CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
86+
isOverride = true
87+
)
88+
}
89+
8190
delay(FEATURE_CONFIG_POLL_INTERVAL_IN_MS)
8291
}
8392
}

‎plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,9 +143,114 @@ class CodeWhispererModelConfiguratorTest {
143143
FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("overrideArn").build())
144144
)
145145
}
146+
abManager.getCustomizationFeature()?.let { customization ->
147+
sut.switchCustomization(
148+
projectRule.project,
149+
CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
150+
isOverride = true
151+
)
152+
}
146153
assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("overrideArn", "foo", null))
147154
}
148155

156+
@Test
157+
fun `should update customization when user has never selected one`() {
158+
val ssoConn = LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url-1", scopes = Q_SCOPES)
159+
ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn)
160+
161+
// Step 1: Server pushes first customization (arnOverride1)
162+
abManager.stub {
163+
on { getCustomizationFeature() }.thenReturn(
164+
FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("arnOverride1").build())
165+
)
166+
}
167+
abManager.getCustomizationFeature()?.let { customization ->
168+
sut.switchCustomization(
169+
projectRule.project,
170+
CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
171+
isOverride = true
172+
)
173+
}
174+
175+
// User should receive arnOverride1 from the server
176+
assertThat(sut.activeCustomization(projectRule.project))
177+
.isEqualTo(CodeWhispererCustomization("arnOverride1", "foo", null))
178+
179+
// Step 2: Server updates customization (arnOverride2)
180+
abManager.stub {
181+
on { getCustomizationFeature() }.thenReturn(
182+
FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("arnOverride2").build())
183+
)
184+
}
185+
186+
abManager.getCustomizationFeature()?.let { customization ->
187+
sut.switchCustomization(
188+
projectRule.project,
189+
CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
190+
isOverride = true
191+
)
192+
}
193+
// User should receive arnOverride2 from the server
194+
assertThat(sut.activeCustomization(projectRule.project))
195+
.isEqualTo(CodeWhispererCustomization("arnOverride2", "foo", null))
196+
}
197+
198+
@Test
199+
fun `should not override user selection when server updates customization`() {
200+
val ssoConn = LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url-1", scopes = Q_SCOPES)
201+
ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn)
202+
203+
// Step 1: Server pushes first customization (arnOverride1)
204+
abManager.stub {
205+
on { getCustomizationFeature() }.thenReturn(
206+
FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("arnOverride1").build())
207+
)
208+
}
209+
210+
abManager.getCustomizationFeature()?.let { customization ->
211+
sut.switchCustomization(
212+
projectRule.project,
213+
CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
214+
isOverride = true
215+
)
216+
}
217+
// User should receive arnOverride1 from the server
218+
assertThat(sut.activeCustomization(projectRule.project))
219+
.isEqualTo(CodeWhispererCustomization("arnOverride1", "foo", null))
220+
221+
// Step 2: Server updates customization again (arnOverride2)
222+
abManager.stub {
223+
on { getCustomizationFeature() }.thenReturn(
224+
FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("arnOverride2").build())
225+
)
226+
}
227+
228+
abManager.getCustomizationFeature()?.let { customization ->
229+
sut.switchCustomization(
230+
projectRule.project,
231+
CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
232+
isOverride = true
233+
)
234+
}
235+
// Ensure server’s change is applied
236+
assertThat(sut.activeCustomization(projectRule.project))
237+
.isEqualTo(CodeWhispererCustomization("arnOverride2", "foo", null))
238+
239+
// Step 3: User manually selects a different customization (userSelectedArn)
240+
val userCustomization = CodeWhispererCustomization("userSelectedArn", "userChoice", null)
241+
sut.switchCustomization(projectRule.project, userCustomization)
242+
abManager.getCustomizationFeature()?.let { customization ->
243+
sut.switchCustomization(
244+
projectRule.project,
245+
CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
246+
isOverride = true
247+
)
248+
}
249+
// Ensure user selection is still respected (should not change to arnOverride2)
250+
assertThat(sut.activeCustomization(projectRule.project))
251+
.isEqualTo(userCustomization)
252+
}
253+
149254
@Test
150255
fun `loadState should load the correct values into memory`() {
151256
credManager.clear()
@@ -315,6 +420,8 @@ class CodeWhispererModelConfiguratorTest {
315420
"fake-sso-url" to CodeWhispererCustomization(arn = "arn_2", name = "name_2", description = "description_2")
316421
)
317422
)
423+
424+
this.serviceDefaultArn = "arn:aws:codewhisperer:default"
318425
}
319426

320427
XmlSerializer.serializeInto(state, element)
@@ -346,6 +453,7 @@ class CodeWhispererModelConfiguratorTest {
346453
"</entry>" +
347454
"</map>" +
348455
"</option>" +
456+
"<option name=\"serviceDefaultArn\" value=\"arn:aws:codewhisperer:default\" />" +
349457
"</component>"
350458

351459
assertThat(actual).isEqualTo(expected)
@@ -362,6 +470,7 @@ class CodeWhispererModelConfiguratorTest {
362470
val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java)
363471
assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(0)
364472
assertThat(actual.previousAvailableCustomizations).hasSize(0)
473+
assertThat(actual.serviceDefaultArn).isNull()
365474
}
366475

367476
@Test
@@ -395,10 +504,12 @@ class CodeWhispererModelConfiguratorTest {
395504
</entry>
396505
</map>
397506
</option>
507+
<option name="serviceDefaultArn" value="arn:aws:codewhisperer:default"/>
398508
</component>
399509
"""
400510
)
401511
val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java)
512+
assertThat(actual.serviceDefaultArn).isEqualTo("arn:aws:codewhisperer:default")
402513
assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(1)
403514
assertThat(actual.connectionIdToActiveCustomizationArn["fake-sso-url"]).isEqualTo(
404515
CodeWhispererCustomization(
@@ -435,6 +546,7 @@ class CodeWhispererModelConfiguratorTest {
435546
)
436547
val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java)
437548
assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(0)
549+
assertThat(actual.serviceDefaultArn).isNull()
438550
assertThat(actual.previousAvailableCustomizations).hasSize(1)
439551
assertThat(actual.previousAvailableCustomizations["fake-sso-url"]).isEqualTo(listOf("arn_1", "arn_2", "arn_3"))
440552
}

0 commit comments

Comments
 (0)