Skip to content

Commit 582fb89

Browse files
Will-ShaoHuaWill-ShaoHua
authored
Merge branch 'main' into project-context
2 parents 5991546 + b72f78b commit 582fb89

File tree

11 files changed

+168
-61
lines changed

11 files changed

+168
-61
lines changed

‎gradle/libs.versions.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ mockitoKotlin = "5.4.0"
2727
mockk = "1.13.10"
2828
nimbus-jose-jwt = "9.40"
2929
node-gradle = "7.0.2"
30-
telemetryGenerator = "1.0.295"
30+
telemetryGenerator = "1.0.297"
3131
testLogger = "4.0.0"
3232
testRetry = "1.5.10"
3333
# test-only; platform provides slf4j transitively at runtime

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,9 @@ class CodeTestChatApp(private val scope: CoroutineScope) : AmazonQApp {
3939
"start-test-gen" to IncomingCodeTestMessage.StartTestGen::class,
4040
"response-body-link-click" to IncomingCodeTestMessage.ClickedLink::class,
4141
"button-click" to IncomingCodeTestMessage.ButtonClicked::class,
42+
"chat-item-voted" to IncomingCodeTestMessage.ChatItemVoted::class,
43+
"chat-item-feedback" to IncomingCodeTestMessage.ChatItemFeedback::class,
44+
"button-click" to IncomingCodeTestMessage.ButtonClicked::class,
4245
"auth-follow-up-was-clicked" to IncomingCodeTestMessage.AuthFollowUpWasClicked::class
4346
)
4447

@@ -80,6 +83,8 @@ class CodeTestChatApp(private val scope: CoroutineScope) : AmazonQApp {
8083
is IncomingCodeTestMessage.StartTestGen -> inboundAppMessagesHandler.processStartTestGen(message)
8184
is IncomingCodeTestMessage.ClickedLink -> inboundAppMessagesHandler.processLinkClick(message)
8285
is IncomingCodeTestMessage.ButtonClicked -> inboundAppMessagesHandler.processButtonClickedMessage(message)
86+
is IncomingCodeTestMessage.ChatItemVoted -> inboundAppMessagesHandler.processChatItemVoted(message)
87+
is IncomingCodeTestMessage.ChatItemFeedback -> inboundAppMessagesHandler.processChatItemFeedBack(message)
8388
is IncomingCodeTestMessage.AuthFollowUpWasClicked -> inboundAppMessagesHandler.processAuthFollowUpClick(message)
8489
}
8590
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,5 +22,9 @@ interface InboundAppMessagesHandler {
2222

2323
suspend fun processButtonClickedMessage(message: IncomingCodeTestMessage.ButtonClicked)
2424

25+
suspend fun processChatItemVoted(message: IncomingCodeTestMessage.ChatItemVoted)
26+
27+
suspend fun processChatItemFeedBack(message: IncomingCodeTestMessage.ChatItemFeedback)
28+
2529
suspend fun processAuthFollowUpClick(message: IncomingCodeTestMessage.AuthFollowUpWasClicked)
2630
}

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

Lines changed: 62 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
// SPDX-License-Identifier: Apache-2.0
33

44
package software.aws.toolkits.jetbrains.services.amazonqCodeTest.controller
5-
5+
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
66
import com.intellij.diff.DiffContentFactory
77
import com.intellij.diff.DiffManager
88
import com.intellij.diff.DiffManagerEx
@@ -44,8 +44,11 @@ import software.amazon.awssdk.services.codewhispererstreaming.model.TextDocument
4444
import software.amazon.awssdk.services.codewhispererstreaming.model.UserInputMessage
4545
import software.amazon.awssdk.services.codewhispererstreaming.model.UserInputMessageContext
4646
import software.amazon.awssdk.services.codewhispererstreaming.model.UserIntent
47+
import software.amazon.awssdk.services.toolkittelemetry.model.Sentiment
4748
import software.aws.toolkits.core.utils.debug
4849
import software.aws.toolkits.core.utils.getLogger
50+
import software.aws.toolkits.core.utils.info
51+
import software.aws.toolkits.core.utils.warn
4952
import software.aws.toolkits.jetbrains.core.AwsClientManager
5053
import software.aws.toolkits.jetbrains.core.coroutines.EDT
5154
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
@@ -81,14 +84,19 @@ import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.TriggerTy
8184
import software.aws.toolkits.jetbrains.services.cwc.clients.chat.v1.ChatSessionV1.Companion.validLanguages
8285
import software.aws.toolkits.jetbrains.services.cwc.controller.chat.StaticPrompt
8386
import software.aws.toolkits.jetbrains.services.cwc.controller.chat.StaticTextResponse
87+
import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.FeedbackComment
8488
import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl
8589
import software.aws.toolkits.jetbrains.services.cwc.editor.context.ActiveFileContext
8690
import software.aws.toolkits.jetbrains.services.cwc.editor.context.ActiveFileContextExtractor
8791
import software.aws.toolkits.jetbrains.services.cwc.editor.context.ExtractionTriggerType
8892
import software.aws.toolkits.jetbrains.services.cwc.editor.context.file.FileContext
8993
import software.aws.toolkits.jetbrains.services.cwc.messages.ChatMessageType
94+
import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService
95+
import software.aws.toolkits.jetbrains.utils.notifyError
9096
import software.aws.toolkits.resources.message
9197
import software.aws.toolkits.telemetry.AmazonqTelemetry
98+
import software.aws.toolkits.telemetry.FeatureId
99+
import software.aws.toolkits.telemetry.InteractionType
92100
import software.aws.toolkits.telemetry.MetricResult
93101
import software.aws.toolkits.telemetry.UiTelemetry
94102
import java.io.File
@@ -403,6 +411,57 @@ class CodeTestChatController(
403411
.build()
404412
}
405413

414+
override suspend fun processChatItemFeedBack(message: IncomingCodeTestMessage.ChatItemFeedback) {
415+
LOG.debug { "$FEATURE_NAME: Processing ChatItemFeedBackMessage: ${message.comment}" }
416+
417+
val session = codeTestChatHelper.getActiveSession()
418+
419+
val comment = FeedbackComment(
420+
conversationId = session.startTestGenerationRequestId,
421+
userComment = message.comment.orEmpty(),
422+
reason = message.selectedOption,
423+
type = "testgen-chat-answer-feedback",
424+
messageId = "",
425+
)
426+
427+
try {
428+
TelemetryService.getInstance().sendFeedback(
429+
sentiment = Sentiment.NEGATIVE,
430+
comment = objectMapper.writeValueAsString(comment),
431+
)
432+
LOG.info { "$FEATURE_NAME answer feedback sent: \"Negative\"" }
433+
} catch (e: Throwable) {
434+
e.notifyError(message("feedback.submit_failed", e))
435+
LOG.warn(e) { "Failed to submit feedback" }
436+
return
437+
}
438+
}
439+
440+
override suspend fun processChatItemVoted(message: IncomingCodeTestMessage.ChatItemVoted) {
441+
LOG.debug { "$FEATURE_NAME: Processing ChatItemVotedMessage: $message" }
442+
443+
val session = codeTestChatHelper.getActiveSession()
444+
when (message.vote) {
445+
"upvote" -> {
446+
AmazonqTelemetry.feedback(
447+
featureId = FeatureId.AmazonQTest,
448+
interactionType = InteractionType.Upvote,
449+
credentialStartUrl = getStartUrl(project = context.project),
450+
amazonqConversationId = session.startTestGenerationRequestId
451+
452+
)
453+
}
454+
"downvote" -> {
455+
AmazonqTelemetry.feedback(
456+
featureId = FeatureId.AmazonQTest,
457+
interactionType = InteractionType.Downvote,
458+
credentialStartUrl = getStartUrl(project = context.project),
459+
amazonqConversationId = session.startTestGenerationRequestId
460+
)
461+
}
462+
}
463+
}
464+
406465
override suspend fun processNewTabCreatedMessage(message: IncomingCodeTestMessage.NewTabCreated) {
407466
newTabOpened(message.tabId)
408467
LOG.debug { "$FEATURE_NAME: New tab created: $message" }
@@ -1312,5 +1371,7 @@ class CodeTestChatController(
13121371

13131372
companion object {
13141373
private val LOG = getLogger<CodeTestChatController>()
1374+
1375+
private val objectMapper = jacksonObjectMapper()
13151376
}
13161377
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,17 @@ sealed interface IncomingCodeTestMessage : CodeTestBaseMessage {
9595
@JsonProperty("actionID") val actionID: String,
9696
) : IncomingCodeTestMessage
9797

98+
data class ChatItemVoted(
99+
@JsonProperty("tabID") val tabId: String,
100+
val vote: String,
101+
) : IncomingCodeTestMessage
102+
103+
data class ChatItemFeedback(
104+
@JsonProperty("tabID") val tabId: String,
105+
val selectedOption: String,
106+
val comment: String?,
107+
) : IncomingCodeTestMessage
108+
98109
data class AuthFollowUpWasClicked(
99110
@JsonProperty("tabID") val tabId: String,
100111
val authType: AuthFollowUpType,

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/chat/telemetry/TelemetryHelper.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -444,7 +444,7 @@ class TelemetryHelper(private val project: Project, private val sessionStorage:
444444

445445
data class FeedbackComment(
446446
val conversationId: String,
447-
val messageId: String,
447+
val messageId: String?,
448448
val reason: String,
449449
val userComment: String,
450450
val type: String = "codewhisperer-chat-answer-feedback",

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt‎

Lines changed: 55 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package software.aws.toolkits.jetbrains.services.codewhisperer.credentials
55

66
import com.intellij.openapi.Disposable
7+
import com.intellij.openapi.application.ApplicationManager
78
import com.intellij.openapi.components.service
89
import com.intellij.openapi.project.Project
910
import com.intellij.util.text.nullize
@@ -40,10 +41,14 @@ import software.amazon.awssdk.services.codewhispererruntime.model.TargetCode
4041
import software.amazon.awssdk.services.codewhispererruntime.model.UserIntent
4142
import software.aws.toolkits.core.utils.debug
4243
import software.aws.toolkits.core.utils.getLogger
44+
import software.aws.toolkits.core.utils.warn
4345
import software.aws.toolkits.jetbrains.core.AwsClientManager
4446
import software.aws.toolkits.jetbrains.core.awsClient
47+
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
48+
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
4549
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
46-
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
50+
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
51+
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
4752
import software.aws.toolkits.jetbrains.services.amazonq.codeWhispererUserContext
4853
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
4954
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
@@ -62,6 +67,7 @@ import java.util.concurrent.TimeUnit
6267
import kotlin.reflect.KProperty0
6368
import kotlin.reflect.jvm.isAccessible
6469

70+
// TODO: move this file to package "/client"
6571
// As the connection is project-level, we need to make this project-level too
6672
@Deprecated("Methods can throw a NullPointerException if callee does not check if connection is valid")
6773
interface CodeWhispererClientAdaptor : Disposable {
@@ -277,16 +283,37 @@ interface CodeWhispererClientAdaptor : Disposable {
277283
open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispererClientAdaptor {
278284
private val mySigv4Client by lazy { createUnmanagedSigv4Client() }
279285

286+
@Volatile
287+
private var myBearerClient: CodeWhispererRuntimeClient? = null
288+
280289
private val KProperty0<*>.isLazyInitialized: Boolean
281290
get() {
282291
isAccessible = true
283292
return (getDelegate() as Lazy<*>).isInitialized()
284293
}
285294

286-
fun bearerClient(): CodeWhispererRuntimeClient =
287-
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.getConnectionSettings()
288-
?.awsClient<CodeWhispererRuntimeClient>()
289-
?: throw Exception("attempt to get bearer client while there is no valid credential")
295+
init {
296+
initClientUpdateListener()
297+
}
298+
299+
private fun initClientUpdateListener() {
300+
ApplicationManager.getApplication().messageBus.connect(this).subscribe(
301+
ToolkitConnectionManagerListener.TOPIC,
302+
object : ToolkitConnectionManagerListener {
303+
override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
304+
if (newConnection is AwsBearerTokenConnection) {
305+
myBearerClient = getBearerClient(newConnection.getConnectionSettings().providerId)
306+
}
307+
}
308+
}
309+
)
310+
}
311+
312+
private fun bearerClient(): CodeWhispererRuntimeClient {
313+
if (myBearerClient != null) return myBearerClient as CodeWhispererRuntimeClient
314+
myBearerClient = getBearerClient()
315+
return myBearerClient as CodeWhispererRuntimeClient
316+
}
290317

291318
override fun generateCompletionsPaginator(firstRequest: GenerateCompletionsRequest) = sequence<GenerateCompletionsResponse> {
292319
var nextToken: String? = firstRequest.nextToken()
@@ -827,6 +854,28 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
827854
if (this::mySigv4Client.isLazyInitialized) {
828855
mySigv4Client.close()
829856
}
857+
myBearerClient?.close()
858+
}
859+
860+
/**
861+
* Every different SSO/AWS Builder ID connection requires a new client which has its corresponding bearer token provider,
862+
* thus we have to create them dynamically.
863+
* Invalidate and recycle the old client first, and create a new client with the new connection.
864+
* This makes sure when we invoke CW, we always use the up-to-date connection.
865+
* In case this fails to close the client, myBearerClient is already set to null thus next time when we invoke CW,
866+
* it will go through this again which should get the current up-to-date connection. This stale client would be
867+
* unused and stay in memory for a while until eventually closed by ToolkitClientManager.
868+
*/
869+
open fun getBearerClient(oldProviderIdToRemove: String = ""): CodeWhispererRuntimeClient? {
870+
myBearerClient = null
871+
872+
val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())
873+
connection as? AwsBearerTokenConnection ?: run {
874+
LOG.warn { "$connection is not a bearer token connection" }
875+
return null
876+
}
877+
878+
return AwsClientManager.getInstance().getClient<CodeWhispererRuntimeClient>(connection.getConnectionSettings())
830879
}
831880

832881
companion object {
@@ -840,6 +889,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
840889
}
841890

842891
class MockCodeWhispererClientAdaptor(override val project: Project) : CodeWhispererClientAdaptorImpl(project) {
892+
override fun getBearerClient(oldProviderIdToRemove: String): CodeWhispererRuntimeClient = project.awsClient()
843893
override fun dispose() {}
844894
}
845895

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

Lines changed: 2 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ import org.junit.After
1515
import org.junit.Before
1616
import org.junit.Rule
1717
import org.junit.Test
18-
import org.junit.jupiter.api.assertThrows
1918
import org.mockito.kotlin.any
2019
import org.mockito.kotlin.argThat
2120
import org.mockito.kotlin.argumentCaptor
@@ -69,20 +68,17 @@ import software.aws.toolkits.core.TokenConnectionSettings
6968
import software.aws.toolkits.core.utils.test.aString
7069
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
7170
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
72-
import software.aws.toolkits.jetbrains.core.credentials.DefaultToolkitConnectionManager
7371
import software.aws.toolkits.jetbrains.core.credentials.ManagedSsoProfile
7472
import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRule
7573
import software.aws.toolkits.jetbrains.core.credentials.MockToolkitAuthManagerRule
7674
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
77-
import software.aws.toolkits.jetbrains.core.credentials.logoutFromSsoConnection
78-
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
79-
import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES
8075
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_REGION
8176
import software.aws.toolkits.jetbrains.services.amazonq.FEATURE_EVALUATION_PRODUCT_NAME
8277
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata
8378
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonRequest
8479
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponseWithToken
8580
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.sdkHttpResponse
81+
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
8682
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptorImpl
8783
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
8884
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
@@ -113,7 +109,7 @@ class CodeWhispererClientAdaptorTest {
113109
private lateinit var bearerClient: CodeWhispererRuntimeClient
114110
private lateinit var ssoClient: SsoOidcClient
115111

116-
private lateinit var sut: CodeWhispererClientAdaptorImpl
112+
private lateinit var sut: CodeWhispererClientAdaptor
117113
private lateinit var connectionManager: ToolkitConnectionManager
118114
private var isTelemetryEnabledDefault: Boolean = false
119115

@@ -167,41 +163,6 @@ class CodeWhispererClientAdaptorTest {
167163
assertThat("us-east-1").isEqualTo(SONO_REGION)
168164
}
169165

170-
@Test
171-
fun `should throw if there is no valid credential, otherwise return codewhispererRuntimeClient`() {
172-
val connectionManager = DefaultToolkitConnectionManager()
173-
projectRule.project.replaceService(ToolkitConnectionManager::class.java, DefaultToolkitConnectionManager(), disposableRule.disposable)
174-
175-
assertThat(ToolkitConnectionManager.getInstance(projectRule.project).activeConnectionForFeature(QConnection.getInstance())).isNull()
176-
assertThrows<Exception>("attempt to get bearer client while there is no valid credential") {
177-
sut.bearerClient()
178-
}
179-
180-
val qConnection = authManagerRule.createConnection(ManagedSsoProfile("us-east-1", aString(), Q_SCOPES))
181-
connectionManager.switchConnection(qConnection)
182-
assertThat(connectionManager.activeConnectionForFeature(QConnection.getInstance()))
183-
.isNotNull
184-
.isEqualTo(qConnection)
185-
assertThat(sut.bearerClient())
186-
.isNotNull
187-
.isInstanceOf(CodeWhispererRuntimeClient::class.java)
188-
189-
logoutFromSsoConnection(projectRule.project, qConnection as AwsBearerTokenConnection)
190-
assertThat(connectionManager.activeConnectionForFeature(QConnection.getInstance())).isNull()
191-
assertThrows<Exception>("attempt to get bearer client while there is no valid credential") {
192-
sut.bearerClient()
193-
}
194-
195-
val anotherQConnection = authManagerRule.createConnection(ManagedSsoProfile("us-east-1", aString(), Q_SCOPES))
196-
connectionManager.switchConnection(anotherQConnection)
197-
assertThat(connectionManager.activeConnectionForFeature(QConnection.getInstance()))
198-
.isNotNull
199-
.isEqualTo(anotherQConnection)
200-
assertThat(sut.bearerClient())
201-
.isNotNull
202-
.isInstanceOf(CodeWhispererRuntimeClient::class.java)
203-
}
204-
205166
@Test
206167
fun `listCustomizations`() {
207168
val sdkIterable = ListAvailableCustomizationsIterable(bearerClient, ListAvailableCustomizationsRequest.builder().build())

0 commit comments

Comments
 (0)