Correcting auth checks for /test workflow

Author: ashishrp-aws

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

@@ -38,7 +38,8 @@ class CodeTestChatApp(private val scope: CoroutineScope) : AmazonQApp {
             "tab-was-removed" to IncomingCodeTestMessage.TabRemoved::class,
             "start-test-gen" to IncomingCodeTestMessage.StartTestGen::class,
             "response-body-link-click" to IncomingCodeTestMessage.ClickedLink::class,
-            "button-click" to IncomingCodeTestMessage.ButtonClicked::class
+            "button-click" to IncomingCodeTestMessage.ButtonClicked::class,
+            "auth-follow-up-was-clicked" to IncomingCodeTestMessage.AuthFollowUpWasClicked::class
         )
 
         scope.launch {
@@ -79,6 +80,7 @@ class CodeTestChatApp(private val scope: CoroutineScope) : AmazonQApp {
             is IncomingCodeTestMessage.StartTestGen -> inboundAppMessagesHandler.processStartTestGen(message)
             is IncomingCodeTestMessage.ClickedLink -> inboundAppMessagesHandler.processLinkClick(message)
             is IncomingCodeTestMessage.ButtonClicked -> inboundAppMessagesHandler.processButtonClickedMessage(message)
+            is IncomingCodeTestMessage.AuthFollowUpWasClicked -> inboundAppMessagesHandler.processAuthFollowUpClick(message)
         }
     }
 

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

@@ -21,4 +21,6 @@ interface InboundAppMessagesHandler {
     suspend fun processTabRemovedMessage(message: IncomingCodeTestMessage.TabRemoved)
 
     suspend fun processButtonClickedMessage(message: IncomingCodeTestMessage.ButtonClicked)
+
+    suspend fun processAuthFollowUpClick(message: IncomingCodeTestMessage.AuthFollowUpWasClicked)
 }

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

@@ -946,6 +946,13 @@ class CodeTestChatController(
         }
     }
 
+    override suspend fun processAuthFollowUpClick(message: IncomingCodeTestMessage.AuthFollowUpWasClicked) {
+        codeTestChatHelper.sendUpdatePromptProgress(message.tabId, null)
+        authController.handleAuth(context.project, message.authType)
+        codeTestChatHelper.sendAuthenticationInProgressMessage(message.tabId) // show user that authentication is in progress
+        codeTestChatHelper.sendChatInputEnabledMessage(false) // disable the input field while authentication is in progress
+    }
+
     private suspend fun updateBuildAndExecuteProgressCard(
         currentStatus: BuildAndExecuteProgressStatus,
         messageId: String?,
@@ -1202,6 +1209,17 @@ class CodeTestChatController(
      * */
     private suspend fun handleChat(tabId: String, message: String) {
         val session = codeTestChatHelper.getActiveSession()
+        session.projectRoot
+        val credentialState = authController.getAuthNeededStates(context.project).amazonQ
+        if (credentialState != null) {
+            messenger.sendAuthNeededException(
+                tabId = tabId,
+                triggerId = UUID.randomUUID().toString(),
+                credentialState = credentialState,
+            )
+            session.isAuthenticating = true
+            return
+        }
         LOG.debug {
             "$FEATURE_NAME: " +
                 "Processing message: $message " +

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

@@ -4,6 +4,7 @@
 package software.aws.toolkits.jetbrains.services.amazonqCodeTest.controller
 
 import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
+import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.ChatInputEnabledMessage
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.CodeTestAddAnswerMessage
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.CodeTestChatMessage
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.CodeTestChatMessageContent
@@ -15,6 +16,8 @@ import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.UpdateP
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.Session
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.storage.ChatSessionStorage
 import software.aws.toolkits.jetbrains.services.cwc.messages.ChatMessageType
+import software.aws.toolkits.jetbrains.services.cwc.messages.FollowUp
+import software.aws.toolkits.resources.message
 import java.util.UUID
 
 class CodeTestChatHelper(
@@ -153,4 +156,27 @@ class CodeTestChatHelper(
             )
         )
     }
+
+    suspend fun sendChatInputEnabledMessage(isEnabled: Boolean) {
+        if (isInvalidSession()) return
+        messagePublisher.publish(ChatInputEnabledMessage(activeCodeTestTabId as String, enabled = isEnabled))
+    }
+
+    suspend fun sendAuthenticationInProgressMessage(
+        tabId: String,
+        messageId: String? = null,
+        followUp: List<FollowUp>? = null,
+        canBeVoted: Boolean? = false,
+    ) {
+        val chatMessage =
+            CodeTestChatMessage(
+                tabId = tabId,
+                messageId = messageId ?: UUID.randomUUID().toString(),
+                messageType = ChatMessageType.Answer,
+                message = message("amazonqFeatureDev.follow_instructions_for_authentication"),
+                followUps = followUp,
+                canBeVoted = canBeVoted ?: false,
+            )
+        messagePublisher.publish(chatMessage)
+    }
 }

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

@@ -94,6 +94,11 @@ sealed interface IncomingCodeTestMessage : CodeTestBaseMessage {
         @JsonProperty("tabID") val tabId: String,
         @JsonProperty("actionID") val actionID: String,
     ) : IncomingCodeTestMessage
+
+    data class AuthFollowUpWasClicked(
+        @JsonProperty("tabID") val tabId: String,
+        val authType: AuthFollowUpType,
+    ) : IncomingCodeTestMessage
 }
 
 data class UpdatePlaceholderMessage(