fix(amazonqFeatureDev): add token for stop gen

Author: tverney

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevApp.kt

@@ -41,7 +41,8 @@ class FeatureDevApp : AmazonQApp {
             "response-body-link-click" to IncomingFeatureDevMessage.ClickedLink::class,
             "insert_code_at_cursor_position" to IncomingFeatureDevMessage.InsertCodeAtCursorPosition::class,
             "open-diff" to IncomingFeatureDevMessage.OpenDiff::class,
-            "file-click" to IncomingFeatureDevMessage.FileClicked::class
+            "file-click" to IncomingFeatureDevMessage.FileClicked::class,
+            "stop-response" to IncomingFeatureDevMessage.StopResponse::class
         )
 
         scope.launch {
@@ -82,6 +83,7 @@ class FeatureDevApp : AmazonQApp {
             is IncomingFeatureDevMessage.InsertCodeAtCursorPosition -> inboundAppMessagesHandler.processInsertCodeAtCursorPosition(message)
             is IncomingFeatureDevMessage.OpenDiff -> inboundAppMessagesHandler.processOpenDiff(message)
             is IncomingFeatureDevMessage.FileClicked -> inboundAppMessagesHandler.processFileClicked(message)
+            is IncomingFeatureDevMessage.StopResponse -> inboundAppMessagesHandler.processStopMessage(message)
         }
     }
 

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

@@ -17,4 +17,5 @@ interface InboundAppMessagesHandler {
     suspend fun processInsertCodeAtCursorPosition(message: IncomingFeatureDevMessage.InsertCodeAtCursorPosition)
     suspend fun processOpenDiff(message: IncomingFeatureDevMessage.OpenDiff)
     suspend fun processFileClicked(message: IncomingFeatureDevMessage.FileClicked)
+    suspend fun processStopMessage(message: IncomingFeatureDevMessage.StopResponse)
 }

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt

@@ -90,6 +90,10 @@ class FeatureDevController(
         )
     }
 
+    override suspend fun processStopMessage(message: IncomingFeatureDevMessage.StopResponse) {
+        handleStopMessage(message)
+    }
+
     override suspend fun processNewTabCreatedMessage(message: IncomingFeatureDevMessage.NewTabCreated) {
         newTabOpened(message.tabId)
     }
@@ -282,6 +286,21 @@ class FeatureDevController(
         }
     }
 
+    private suspend fun handleStopMessage(message: IncomingFeatureDevMessage.StopResponse) {
+        var session: Session? = null
+        messenger.sendAnswer( tabId = message.tabId, "Stopping code generation ...", messageType = FeatureDevMessageType.Answer, canBeVoted = false)
+        messenger.sendUpdatePlaceholder(
+            tabId = message.tabId,
+            newPlaceholder = "Stopping code generation ..."
+        )
+        messenger.sendChatInputEnabledMessage(tabId = message.tabId, enabled = false)
+        session = getSessionInfo(message.tabId)
+
+        if (session.sessionState.token?.token() !== null) {
+            session.sessionState?.token?.cancel()
+        }
+
+    }
     private suspend fun insertCode(tabId: String) {
         var session: Session? = null
         try {

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/messages/FeatureDevMessage.kt

@@ -66,6 +66,12 @@ sealed interface IncomingFeatureDevMessage : FeatureDevBaseMessage {
         val link: String,
     ) : IncomingFeatureDevMessage
 
+    data class StopResponse(
+        @JsonProperty("tabID") val tabId: String,
+        val command: String,
+        val messageId: String?
+    ) : IncomingFeatureDevMessage
+
     data class InsertCodeAtCursorPosition(
         @JsonProperty("tabID") val tabId: String,
         val code: String,

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationState.kt

@@ -4,6 +4,7 @@
 package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session
 
 import kotlinx.coroutines.delay
+import org.gradle.tooling.CancellationTokenSource
 import software.amazon.awssdk.services.codewhispererruntime.model.CodeGenerationWorkflowStatus
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.warn
@@ -24,6 +25,7 @@ private val logger = getLogger<CodeGenerationState>()
 class CodeGenerationState(
     override val tabID: String,
     override var approach: String,
+    override val token: CancellationTokenSource?,
     val config: SessionStateConfig,
     val uploadId: String,
     val currentIteration: Int,
@@ -78,7 +80,8 @@ class CodeGenerationState(
                 uploadId = uploadId,
                 messenger = messenger,
                 codeGenerationRemainingIterationCount = codeGenerationRemainingIterationCount,
-                codeGenerationTotalIterationCount = codeGenerationTotalIterationCount
+                codeGenerationTotalIterationCount = codeGenerationTotalIterationCount,
+                token = this.token
             )
 
             // It is not needed to interact right away with the PrepareCodeGeneration.

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/ConversationNotStartedState.kt

@@ -3,10 +3,13 @@
 
 package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session
 
+import org.gradle.tooling.CancellationTokenSource
+
 class ConversationNotStartedState(
     override var approach: String,
     override val tabID: String,
-) : SessionState {
+    override val token: CancellationTokenSource?,
+    ) : SessionState {
     override val phase = SessionStatePhase.INIT
 
     override suspend fun interact(action: SessionStateAction): SessionStateInteraction {

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationState.kt

@@ -3,6 +3,7 @@
 
 package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session
 
+import org.gradle.tooling.CancellationTokenSource
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
@@ -22,6 +23,7 @@ private val logger = getLogger<PrepareCodeGenerationState>()
 
 class PrepareCodeGenerationState(
     override var tabID: String,
+    override val token: CancellationTokenSource?,
     override var approach: String,
     private var config: SessionStateConfig,
     val filePaths: List<NewFileZipInfo>,
@@ -71,7 +73,8 @@ class PrepareCodeGenerationState(
                 uploadId = this.uploadId,
                 currentIteration = this.currentIteration,
                 repositorySize = zipFileLength.toDouble(),
-                messenger = messenger
+                messenger = messenger,
+                token = this.token
             )
         } catch (e: Exception) {
             result = Result.Failed

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/Session.kt

@@ -41,7 +41,7 @@ class Session(val tabID: String, val project: Project) {
         context = FeatureDevSessionContext(project, MAX_PROJECT_SIZE_BYTES)
         proxyClient = FeatureDevClient.getInstance(project)
         featureDevService = FeatureDevService(proxyClient, project)
-        _state = ConversationNotStartedState("", tabID)
+        _state = ConversationNotStartedState("", tabID, null)
         isAuthenticating = false
         codegenRetries = CODE_GENERATION_RETRY_LIMIT
     }
@@ -81,6 +81,7 @@ class Session(val tabID: String, val project: Project) {
             currentIteration = 0, // first code gen iteration
             uploadId = "", // There is no code gen uploadId so far
             messenger = messenger,
+            token = null
         )
     }
 

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/SessionState.kt

@@ -3,9 +3,12 @@
 
 package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session
 
+import org.gradle.tooling.CancellationTokenSource
+
 interface SessionState {
     val tabID: String
     val phase: SessionStatePhase?
+    val token: CancellationTokenSource?
     var approach: String
     suspend fun interact(action: SessionStateAction): SessionStateInteraction
 }

plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/model/BuildStepStatus.kt

@@ -17,4 +17,5 @@ fun mapTransformationPlanApiStatus(apiStatus: TransformationProgressUpdateStatus
     TransformationProgressUpdateStatus.FAILED, TransformationProgressUpdateStatus.PAUSED -> BuildStepStatus.WARNING
     TransformationProgressUpdateStatus.IN_PROGRESS -> BuildStepStatus.WORKING
     TransformationProgressUpdateStatus.UNKNOWN_TO_SDK_VERSION -> BuildStepStatus.ERROR
+    TransformationProgressUpdateStatus.AWAITING_CLIENT_ACTION -> TODO()
 }