Skip to content

Commit 5af8d6c

Browse files
rlirli
authored
Merge pull request #5676 from aws/autoMerge/feature/q-lsp-chat
Merge main into feature/q-lsp-chat
2 parents 04e11d6 + 57c52d4 commit 5af8d6c

File tree

13 files changed

+212
-122
lines changed

13 files changed

+212
-122
lines changed

‎.changes/next-release/feature-ae7d15f4-ae07-4faf-963a-6d83a85852e4.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "feature",
3+
"description" : "Amazon Q: Support selecting customizations across all Q profiles with automatic profile switching for enterprise users"
4+
}

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

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ 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
1716
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
1817
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
1918
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isUserBuilderId
@@ -79,18 +78,16 @@ class CodeWhispererProjectStartupActivity : StartupActivity.DumbAware {
7978
projectCoroutineScope(project).launch {
8079
while (isActive) {
8180
CodeWhispererFeatureConfigService.getInstance().fetchFeatureConfigs(project)
82-
CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.let { customization ->
81+
CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.let { overrideContext ->
8382
val persistedCustomizationOverride = CodeWhispererModelConfigurator.getInstance().getPersistedCustomizationOverride()
84-
val latestCustomizationOverride = customization.value.stringValue()
83+
val latestCustomizationOverride = overrideContext.value.stringValue()
8584
if (persistedCustomizationOverride == latestCustomizationOverride) return@let
8685

8786
// latest is different from the currently persisted, need update
88-
CodeWhispererFeatureConfigService.getInstance().validateCustomizationOverride(project, customization)
89-
CodeWhispererModelConfigurator.getInstance().switchCustomization(
90-
project,
91-
CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
92-
isOverride = true
93-
)
87+
val customization = CodeWhispererFeatureConfigService.getInstance().validateCustomizationOverride(project, overrideContext)
88+
if (customization != null) {
89+
CodeWhispererModelConfigurator.getInstance().switchCustomization(project, customization, isOverride = true)
90+
}
9491
}
9592

9693
delay(FEATURE_CONFIG_POLL_INTERVAL_IN_MS)

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ import org.mockito.kotlin.argumentCaptor
1818
import org.mockito.kotlin.doReturn
1919
import org.mockito.kotlin.mock
2020
import org.mockito.kotlin.stub
21-
import org.mockito.kotlin.times
2221
import org.mockito.kotlin.verify
2322
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
2423
import software.amazon.awssdk.services.codewhispererruntime.model.ArtifactType

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
3535
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
3636
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_URL
3737
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
38+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
3839
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
3940
import kotlin.reflect.full.memberFunctions
4041
import kotlin.test.Test
@@ -78,7 +79,7 @@ class CodeWhispererFeatureConfigServiceTest {
7879

7980
projectRule.project.replaceService(
8081
QRegionProfileManager::class.java,
81-
mock<QRegionProfileManager> { on { getQClient(any<Project>(), eq(CodeWhispererRuntimeClient::class)) } doReturn mockClient },
82+
mock<QRegionProfileManager> { on { getQClient(any<Project>(), eq(QRegionProfile()), eq(CodeWhispererRuntimeClient::class)) } doReturn mockClient },
8283
disposableRule.disposable
8384
)
8485

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

Lines changed: 65 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import migration.software.aws.toolkits.jetbrains.services.codewhisperer.customiz
1414
import org.assertj.core.api.Assertions.assertThat
1515
import org.jdom.output.XMLOutputter
1616
import org.junit.Before
17+
import org.junit.Ignore
1718
import org.junit.Rule
1819
import org.junit.Test
1920
import org.mockito.kotlin.any
@@ -42,6 +43,8 @@ import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
4243
import software.aws.toolkits.jetbrains.core.region.MockRegionProviderRule
4344
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
4445
import software.aws.toolkits.jetbrains.services.amazonq.FeatureContext
46+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
47+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
4548
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
4649
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
4750
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
@@ -82,6 +85,7 @@ class CodeWhispererModelConfiguratorTest {
8285
private lateinit var mockClient: CodeWhispererRuntimeClient
8386
private lateinit var abManager: CodeWhispererFeatureConfigService
8487
private lateinit var mockClintAdaptor: CodeWhispererClientAdaptor
88+
private lateinit var mockQRegionProfileManager: QRegionProfileManager
8589

8690
@Before
8791
fun setup() {
@@ -124,6 +128,11 @@ class CodeWhispererModelConfiguratorTest {
124128

125129
mockClintAdaptor = mock()
126130
projectRule.project.registerServiceInstance(CodeWhispererClientAdaptor::class.java, mockClintAdaptor)
131+
132+
mockQRegionProfileManager = mock {
133+
on { listRegionProfiles(any()) }.thenReturn(listOf(QRegionProfile("fake_name", "fake_arn")))
134+
}
135+
ApplicationManager.getApplication().replaceService(QRegionProfileManager::class.java, mockQRegionProfileManager, disposableRule.disposable)
127136
}
128137

129138
@Test
@@ -431,7 +440,10 @@ class CodeWhispererModelConfiguratorTest {
431440

432441
this.connectionIdToActiveCustomizationArn.putAll(
433442
mapOf(
434-
"fake-sso-url" to CodeWhispererCustomization(arn = "arn_2", name = "name_2", description = "description_2")
443+
"fake-sso-url" to CodeWhispererCustomization(
444+
arn = "arn_2", name = "name_2", description = "description_2",
445+
profile = QRegionProfile(profileName = "myActiveProfile", arn = "arn:aws:codewhisperer:us-west-2:123456789012:profile/myActiveProfile")
446+
)
435447
)
436448
)
437449

@@ -450,6 +462,12 @@ class CodeWhispererModelConfiguratorTest {
450462
"<option name=\"arn\" value=\"arn_2\" />" +
451463
"<option name=\"name\" value=\"name_2\" />" +
452464
"<option name=\"description\" value=\"description_2\" />" +
465+
"<option name=\"profile\">" +
466+
"<QRegionProfile>" +
467+
"<option name=\"arn\" value=\"arn:aws:codewhisperer:us-west-2:123456789012:profile/myActiveProfile\" />" +
468+
"<option name=\"profileName\" value=\"myActiveProfile\" />" +
469+
"</QRegionProfile>" +
470+
"</option>" +
453471
"</CodeWhispererCustomization>" +
454472
"</value>" +
455473
"</entry>" +
@@ -500,6 +518,10 @@ class CodeWhispererModelConfiguratorTest {
500518
<option name="arn" value="arn_2" />
501519
<option name="name" value="name_2" />
502520
<option name="description" value="description_2" />
521+
<option name="profile"><QRegionProfile>
522+
<option name="profileName" value="myActiveProfile"/>
523+
<option name="arn" value="arn:aws:codewhisperer:us-west-2:123456789012:profile/myActiveProfile"/>
524+
</QRegionProfile></option>
503525
</CodeWhispererCustomization>
504526
</value>
505527
</entry>
@@ -529,14 +551,42 @@ class CodeWhispererModelConfiguratorTest {
529551
CodeWhispererCustomization(
530552
arn = "arn_2",
531553
name = "name_2",
532-
description = "description_2"
554+
description = "description_2",
555+
profile = QRegionProfile("myActiveProfile", "arn:aws:codewhisperer:us-west-2:123456789012:profile/myActiveProfile")
533556
)
534557
)
535558

536559
assertThat(actual.previousAvailableCustomizations).hasSize(1)
537560
assertThat(actual.previousAvailableCustomizations["fake-sso-url"]).isEqualTo(listOf("arn_1", "arn_2", "arn_3"))
538561
}
539562

563+
@Test
564+
fun `backward compatibility - should still be deseriealizable where profile field is not present`() {
565+
val xml = xmlElement(
566+
"""
567+
<component name="codewhispererCustomizationStates">
568+
<option name="connectionIdToActiveCustomizationArn">
569+
<map>
570+
<entry key="sso-session:foo">
571+
<value>
572+
<CodeWhispererCustomization>
573+
<option name="arn" value="arn:foo" />
574+
<option name="name" value="Customization-foo" />
575+
<option name="description" value="Foo foo foo foo" />
576+
</CodeWhispererCustomization>
577+
</value>
578+
</entry>
579+
</map>
580+
</option>
581+
</component>
582+
""".trimIndent()
583+
)
584+
585+
val actual = XmlSerializer.deserialize(xml, CodeWhispererCustomizationState::class.java)
586+
val cnt = actual.connectionIdToActiveCustomizationArn.size
587+
assertThat(cnt).isEqualTo(1)
588+
}
589+
540590
@Test
541591
fun `deserialize users choosing default customization`() {
542592
val element = xmlElement(
@@ -565,6 +615,7 @@ class CodeWhispererModelConfiguratorTest {
565615
assertThat(actual.previousAvailableCustomizations["fake-sso-url"]).isEqualTo(listOf("arn_1", "arn_2", "arn_3"))
566616
}
567617

618+
@Ignore
568619
@Test
569620
fun `profile switch should keep using existing customization if new list still contains that arn`() {
570621
val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url 1", scopes = Q_SCOPES))
@@ -573,6 +624,11 @@ class CodeWhispererModelConfiguratorTest {
573624
sut.switchCustomization(projectRule.project, oldCustomization)
574625

575626
assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
627+
// TODO: mock sdk client to fix the test
628+
// val fakeCustomizations = listOf(
629+
// CodeWhispererCustomization("oldArn", "oldName", "oldDescription")
630+
// )
631+
// mockClintAdaptor.stub { on { listAvailableCustomizations(QRegionProfile("fake_name", "fake_arn")) } doReturn fakeCustomizations }
576632

577633
ApplicationManager.getApplication().messageBus
578634
.syncPublisher(QRegionProfileSelectedListener.TOPIC)
@@ -581,6 +637,7 @@ class CodeWhispererModelConfiguratorTest {
581637
assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
582638
}
583639

640+
@Ignore
584641
@Test
585642
fun `profile switch should invalidate obsolete customization if it's not in the new list`() {
586643
val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url 1", scopes = Q_SCOPES))
@@ -589,6 +646,12 @@ class CodeWhispererModelConfiguratorTest {
589646
sut.switchCustomization(projectRule.project, oldCustomization)
590647
assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
591648

649+
// TODO: mock sdk client to fix the test
650+
// val fakeCustomizations = listOf(
651+
// CodeWhispererCustomization("newArn", "newName", "newDescription")
652+
// )
653+
// mockClintAdaptor.stub { on { listAvailableCustomizations(QRegionProfile("fake_name", "fake_arn")) } doReturn fakeCustomizations }
654+
592655
val latch = CountDownLatch(1)
593656

594657
ApplicationManager.getApplication().messageBus

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,7 @@ class QRegionProfileManagerTest {
178178
client.stub {
179179
onGeneric { listAvailableProfilesPaginator(any<Consumer<ListAvailableProfilesRequest.Builder>>()) } doReturn iterable
180180
}
181-
val connectionSettings = sut.getQClientSettings(project)
181+
val connectionSettings = sut.getQClientSettings(project, null)
182182
resourceCache.addEntry(connectionSettings, QProfileResources.LIST_REGION_PROFILES, QProfileResources.LIST_REGION_PROFILES.fetch(connectionSettings))
183183

184184
assertThat(sut.listRegionProfiles(project))
@@ -247,7 +247,7 @@ class QRegionProfileManagerTest {
247247
sut.activeProfile(project)
248248
).isEqualTo(QRegionProfile(arn = "arn:aws:codewhisperer:eu-central-1:123456789012:profile/FOO_PROFILE", profileName = "FOO_PROFILE"))
249249

250-
val settings = sut.getQClientSettings(project)
250+
val settings = sut.getQClientSettings(project, null)
251251
assertThat(settings.region.id).isEqualTo(Region.EU_CENTRAL_1.id())
252252

253253
sut.switchProfile(
@@ -259,7 +259,7 @@ class QRegionProfileManagerTest {
259259
sut.activeProfile(project)
260260
).isEqualTo(QRegionProfile(arn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/BAR_PROFILE", profileName = "BAR_PROFILE"))
261261

262-
val settings2 = sut.getQClientSettings(project)
262+
val settings2 = sut.getQClientSettings(project, null)
263263
assertThat(settings2.region.id).isEqualTo(Region.US_EAST_1.id())
264264
}
265265

@@ -275,7 +275,7 @@ class QRegionProfileManagerTest {
275275
assertThat(
276276
sut.activeProfile(project)
277277
).isEqualTo(QRegionProfile(arn = "arn:aws:codewhisperer:eu-central-1:123456789012:profile/FOO_PROFILE", profileName = "FOO_PROFILE"))
278-
assertThat(sut.getQClientSettings(project).region.id).isEqualTo(Region.EU_CENTRAL_1.id())
278+
assertThat(sut.getQClientSettings(project, null).region.id).isEqualTo(Region.EU_CENTRAL_1.id())
279279

280280
val client = sut.getQClient<CodeWhispererRuntimeClient>(project)
281281
assertThat(client).isInstanceOf(CodeWhispererRuntimeClient::class.java)
@@ -292,7 +292,7 @@ class QRegionProfileManagerTest {
292292
assertThat(
293293
sut.activeProfile(project)
294294
).isEqualTo(QRegionProfile(arn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/BAR_PROFILE", profileName = "BAR_PROFILE"))
295-
assertThat(sut.getQClientSettings(project).region.id).isEqualTo(Region.US_EAST_1.id())
295+
assertThat(sut.getQClientSettings(project, null).region.id).isEqualTo(Region.US_EAST_1.id())
296296

297297
val client2 = sut.getQClient<CodeWhispererRuntimeClient>(project)
298298
assertThat(client2).isInstanceOf(CodeWhispererRuntimeClient::class.java)

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt‎

Lines changed: 38 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import software.aws.toolkits.core.utils.getLogger
1515
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
1616
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
1717
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
18+
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
1819
import software.aws.toolkits.jetbrains.utils.isQExpired
1920

2021
@Service
@@ -103,37 +104,49 @@ class CodeWhispererFeatureConfigService {
103104
}
104105
}
105106

106-
fun validateCustomizationOverride(project: Project, customization: FeatureContext) {
107-
val customizationArnOverride = customization.value.stringValue()
108-
val connection = connection(project) ?: return
109-
if (customizationArnOverride != null) {
110-
// Double check if server-side wrongly returns a customizationArn to BID users
111-
calculateIfBIDConnection(project) {
112-
featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
113-
}
114-
val availableCustomizations =
115-
calculateIfIamIdentityCenterConnection(project) {
116-
try {
107+
fun validateCustomizationOverride(project: Project, featOverrideContext: FeatureContext): CodeWhispererCustomization? {
108+
val customizationArnOverride = featOverrideContext.value.stringValue()
109+
connection(project) ?: return null
110+
customizationArnOverride ?: return null
111+
112+
// Double check if server-side wrongly returns a customizationArn to BID users
113+
calculateIfBIDConnection(project) {
114+
featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
115+
}
116+
val availableCustomizations =
117+
calculateIfIamIdentityCenterConnection(project) {
118+
try {
119+
val profiles = QRegionProfileManager.getInstance().listRegionProfiles(project)
120+
?: error("Attempted to fetch profiles while there does not exist")
121+
122+
val customs = profiles.flatMap { profile ->
117123
QRegionProfileManager.getInstance().getQClient<CodeWhispererRuntimeClient>(project)
118-
.listAvailableCustomizationsPaginator {}
119-
.flatMap { resp ->
120-
resp.customizations().map {
121-
it.arn()
122-
}
124+
.listAvailableCustomizations { it.profileArn(profile.arn) }.customizations().map { originalCustom ->
125+
CodeWhispererCustomization(
126+
arn = originalCustom.arn(),
127+
name = originalCustom.name(),
128+
description = originalCustom.description(),
129+
profile = profile
130+
)
123131
}
124-
} catch (e: Exception) {
125-
LOG.debug(e) { "Failed to list available customizations" }
126-
null
127132
}
128-
}
129133

130-
// If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value
131-
if (availableCustomizations?.contains(customizationArnOverride) == false) {
132-
LOG.debug {
133-
"Customization arn $customizationArnOverride not available in listAvailableCustomizations, not using"
134+
customs
135+
} catch (e: Exception) {
136+
LOG.debug(e) { "encountered error while validating customization override" }
137+
null
134138
}
135-
featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
136139
}
140+
141+
val isValidOverride = availableCustomizations != null && availableCustomizations.any { it.arn == customizationArnOverride }
142+
143+
// If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value
144+
return if (!isValidOverride) {
145+
LOG.debug { "Customization arn $customizationArnOverride not available in listAvailableCustomizations, not using" }
146+
featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
147+
null
148+
} else {
149+
availableCustomizations?.find { it.arn == customizationArnOverride }
137150
}
138151
}
139152

0 commit comments

Comments
 (0)