Skip to content

Commit a4e3442

Browse files
Will-ShaoHuaWill-ShaoHua
committed
Merge remote-tracking branch 'upstream/main' into client
2 parents eab031b + b440147 commit a4e3442

File tree

9 files changed

+202
-11
lines changed

9 files changed

+202
-11
lines changed

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatApp.kt‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,8 @@ class CodeTestChatApp(private val scope: CoroutineScope) : AmazonQApp {
3838
"tab-was-removed" to IncomingCodeTestMessage.TabRemoved::class,
3939
"start-test-gen" to IncomingCodeTestMessage.StartTestGen::class,
4040
"response-body-link-click" to IncomingCodeTestMessage.ClickedLink::class,
41-
"button-click" to IncomingCodeTestMessage.ButtonClicked::class
41+
"button-click" to IncomingCodeTestMessage.ButtonClicked::class,
42+
"auth-follow-up-was-clicked" to IncomingCodeTestMessage.AuthFollowUpWasClicked::class
4243
)
4344

4445
scope.launch {
@@ -79,6 +80,7 @@ class CodeTestChatApp(private val scope: CoroutineScope) : AmazonQApp {
7980
is IncomingCodeTestMessage.StartTestGen -> inboundAppMessagesHandler.processStartTestGen(message)
8081
is IncomingCodeTestMessage.ClickedLink -> inboundAppMessagesHandler.processLinkClick(message)
8182
is IncomingCodeTestMessage.ButtonClicked -> inboundAppMessagesHandler.processButtonClickedMessage(message)
83+
is IncomingCodeTestMessage.AuthFollowUpWasClicked -> inboundAppMessagesHandler.processAuthFollowUpClick(message)
8284
}
8385
}
8486

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/InboundAppMessagesHandler.kt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,4 +21,6 @@ interface InboundAppMessagesHandler {
2121
suspend fun processTabRemovedMessage(message: IncomingCodeTestMessage.TabRemoved)
2222

2323
suspend fun processButtonClickedMessage(message: IncomingCodeTestMessage.ButtonClicked)
24+
25+
suspend fun processAuthFollowUpClick(message: IncomingCodeTestMessage.AuthFollowUpWasClicked)
2426
}

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -954,6 +954,14 @@ class CodeTestChatController(
954954
}
955955
}
956956

957+
override suspend fun processAuthFollowUpClick(message: IncomingCodeTestMessage.AuthFollowUpWasClicked) {
958+
codeTestChatHelper.sendUpdatePromptProgress(message.tabId, null)
959+
authController.handleAuth(context.project, message.authType)
960+
codeTestChatHelper.sendAuthenticationInProgressMessage(message.tabId) // show user that authentication is in progress
961+
codeTestChatHelper.sendChatInputEnabledMessage(false) // disable the input field while authentication is in progress
962+
sessionCleanUp(codeTestChatHelper.getActiveSession().tabId)
963+
}
964+
957965
private suspend fun updateBuildAndExecuteProgressCard(
958966
currentStatus: BuildAndExecuteProgressStatus,
959967
messageId: String?,
@@ -1210,6 +1218,17 @@ class CodeTestChatController(
12101218
* */
12111219
private suspend fun handleChat(tabId: String, message: String) {
12121220
val session = codeTestChatHelper.getActiveSession()
1221+
session.projectRoot
1222+
val credentialState = authController.getAuthNeededStates(context.project).amazonQ
1223+
if (credentialState != null) {
1224+
messenger.sendAuthNeededException(
1225+
tabId = tabId,
1226+
triggerId = UUID.randomUUID().toString(),
1227+
credentialState = credentialState,
1228+
)
1229+
session.isAuthenticating = true
1230+
return
1231+
}
12131232
LOG.debug {
12141233
"$FEATURE_NAME: " +
12151234
"Processing message: $message " +

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatHelper.kt‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package software.aws.toolkits.jetbrains.services.amazonqCodeTest.controller
55

66
import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
7+
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.ChatInputEnabledMessage
78
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.CodeTestAddAnswerMessage
89
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.CodeTestChatMessage
910
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.CodeTestChatMessageContent
@@ -15,6 +16,8 @@ import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.UpdateP
1516
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.Session
1617
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.storage.ChatSessionStorage
1718
import software.aws.toolkits.jetbrains.services.cwc.messages.ChatMessageType
19+
import software.aws.toolkits.jetbrains.services.cwc.messages.FollowUp
20+
import software.aws.toolkits.resources.message
1821
import java.util.UUID
1922

2023
class CodeTestChatHelper(
@@ -153,4 +156,27 @@ class CodeTestChatHelper(
153156
)
154157
)
155158
}
159+
160+
suspend fun sendChatInputEnabledMessage(isEnabled: Boolean) {
161+
if (isInvalidSession()) return
162+
messagePublisher.publish(ChatInputEnabledMessage(activeCodeTestTabId as String, enabled = isEnabled))
163+
}
164+
165+
suspend fun sendAuthenticationInProgressMessage(
166+
tabId: String,
167+
messageId: String? = null,
168+
followUp: List<FollowUp>? = null,
169+
canBeVoted: Boolean? = false,
170+
) {
171+
val chatMessage =
172+
CodeTestChatMessage(
173+
tabId = tabId,
174+
messageId = messageId ?: UUID.randomUUID().toString(),
175+
messageType = ChatMessageType.Answer,
176+
message = message("amazonqFeatureDev.follow_instructions_for_authentication"),
177+
followUps = followUp,
178+
canBeVoted = canBeVoted ?: false,
179+
)
180+
messagePublisher.publish(chatMessage)
181+
}
156182
}

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/messages/CodeTestMessage.kt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,11 @@ sealed interface IncomingCodeTestMessage : CodeTestBaseMessage {
9494
@JsonProperty("tabID") val tabId: String,
9595
@JsonProperty("actionID") val actionID: String,
9696
) : IncomingCodeTestMessage
97+
98+
data class AuthFollowUpWasClicked(
99+
@JsonProperty("tabID") val tabId: String,
100+
val authType: AuthFollowUpType,
101+
) : IncomingCodeTestMessage
97102
}
98103

99104
data class UpdatePlaceholderMessage(

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