Skip to content

Commit efb9a5b

Browse files
ashishrp-awsashishrp-aws
authored
Merge branch 'main' into auth
2 parents becc7a4 + 951d519 commit efb9a5b

File tree

7 files changed

+293
-41
lines changed

7 files changed

+293
-41
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/codetest/sessionconfig/CodeTestSessionConfig.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.noFileOpe
2929
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
3030
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererUnknownLanguage
3131
import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmingLanguage
32+
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
3233
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.CODE_SCAN_CREATE_PAYLOAD_TIMEOUT_IN_SECONDS
3334
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.DEFAULT_CODE_SCAN_TIMEOUT_IN_SECONDS
3435
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.DEFAULT_PAYLOAD_LIMIT_IN_BYTES
@@ -217,7 +218,7 @@ class CodeTestSessionConfig(
217218
for (module in project.modules) {
218219
val changeListManager = ChangeListManager.getInstance(module.project)
219220
module.guessModuleDir()?.let { moduleDir ->
220-
val gitIgnoreFilteringUtil = GitIgnoreFilteringUtil(moduleDir)
221+
val gitIgnoreFilteringUtil = GitIgnoreFilteringUtil(moduleDir, CodeWhispererConstants.FeatureName.TEST_GENERATION)
221222
stack.push(moduleDir)
222223
while (stack.isNotEmpty()) {
223224
val current = stack.pop()

‎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/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/GitIgnoreFilteringUtil.kt‎

Lines changed: 51 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -9,37 +9,58 @@ import kotlinx.coroutines.async
99
import kotlinx.coroutines.withContext
1010
import kotlin.coroutines.coroutineContext
1111

12-
class GitIgnoreFilteringUtil(private val moduleDir: VirtualFile) {
12+
class GitIgnoreFilteringUtil(
13+
private val moduleDir: VirtualFile,
14+
private val useCase: CodeWhispererConstants.FeatureName? = null,
15+
) {
1316
private var ignorePatternsWithGitIgnore = emptyList<Regex>()
14-
private val additionalGitIgnoreRules = setOf(
15-
".aws-sam",
16-
".gem",
17-
".git",
18-
".gitignore",
19-
".gradle",
20-
".hg",
21-
".idea",
22-
".project",
23-
".rvm",
24-
".svn",
25-
"*.zip",
26-
"*.bin",
27-
"*.png",
28-
"*.jpg",
29-
"*.svg",
30-
"*.pyc",
31-
"license.txt",
32-
"License.txt",
33-
"LICENSE.txt",
34-
"license.md",
35-
"License.md",
36-
"LICENSE.md",
37-
"node_modules",
38-
"build",
39-
"dist",
40-
"annotation-generated-src",
41-
"annotation-generated-tst"
42-
)
17+
private val additionalGitIgnoreRules = buildSet {
18+
addAll(
19+
setOf(
20+
".aws-sam",
21+
".gem",
22+
".git",
23+
".gitignore",
24+
".gradle",
25+
".hg",
26+
".idea",
27+
".project",
28+
".rvm",
29+
".svn",
30+
"*.zip",
31+
"*.bin",
32+
"*.png",
33+
"*.jpg",
34+
"*.svg",
35+
"*.pyc",
36+
"license.txt",
37+
"License.txt",
38+
"LICENSE.txt",
39+
"license.md",
40+
"License.md",
41+
"LICENSE.md",
42+
"node_modules",
43+
"build",
44+
"dist",
45+
"annotation-generated-src",
46+
"annotation-generated-tst"
47+
)
48+
)
49+
if (useCase == CodeWhispererConstants.FeatureName.TEST_GENERATION) {
50+
addAll(
51+
setOf(
52+
"env",
53+
"release-info",
54+
"*.jar",
55+
"*.exe",
56+
"*.a",
57+
"*.map",
58+
"*.graph",
59+
"*.so"
60+
)
61+
)
62+
}
63+
}
4364

4465
init {
4566
ignorePatternsWithGitIgnore = try {

‎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)