fix(amazonqFeatureDev): include cancellation token

Author: tverney

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

@@ -288,10 +288,10 @@ 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.sendAnswer( tabId = message.tabId, message("amazonqFeatureDev.code_generation.stopping_code_generation"), messageType = FeatureDevMessageType.Answer, canBeVoted = false)
         messenger.sendUpdatePlaceholder(
             tabId = message.tabId,
-            newPlaceholder = "Stopping code generation ..."
+            newPlaceholder = message("amazonqFeatureDev.code_generation.stopping_code_generation")
         )
         messenger.sendChatInputEnabledMessage(tabId = message.tabId, enabled = false)
         session = getSessionInfo(message.tabId)

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

@@ -4,6 +4,7 @@
 package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.controller
 
 import com.intellij.notification.NotificationAction
+import org.gradle.tooling.GradleConnector
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CODE_GENERATION_RETRY_LIMIT
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FeatureDevMessageType
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUp
@@ -68,6 +69,21 @@ suspend fun FeatureDevController.onCodeGeneration(session: Session, message: Str
             }
         }
 
+        if (state.token?.token()?.isCancellationRequested == true) {
+            messenger.sendAnswer(
+                tabId = tabId,
+                messageType = FeatureDevMessageType.Answer,
+                message = message("amazonqFeatureDev.code_generation.stopped_code_generation")
+            )
+            messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = true)
+
+            messenger.sendUpdatePlaceholder(
+                tabId = tabId,
+                newPlaceholder = message("amazonqFeatureDev.placeholder.new_plan")
+            )
+            return
+        }
+
         // Atm this is the only possible path as codegen is mocked to return empty.
         if (filePaths.size or deletedFiles.size == 0) {
             messenger.sendAnswer(
@@ -115,9 +131,13 @@ suspend fun FeatureDevController.onCodeGeneration(session: Session, message: Str
 
         messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.after_code_generation"))
     } finally {
-        messenger.sendAsyncEventProgress(tabId = tabId, inProgress = false) // Finish processing the event
-        messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false) // Lock chat input until a follow-up is clicked.
 
+        if (session.sessionState.token?.token()?.isCancellationRequested == true) {
+            session.sessionState.token = GradleConnector.newCancellationTokenSource()
+        } else {
+            messenger.sendAsyncEventProgress(tabId = tabId, inProgress = false) // Finish processing the event
+            messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false) // Lock chat input until a follow-up is clicked.
+        }
         if (toolWindow != null && !toolWindow.isVisible) {
             notifyInfo(
                 title = message("amazonqFeatureDev.code_generation.notification_title"),

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

@@ -25,16 +25,16 @@ 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,
     val repositorySize: Double,
     val messenger: MessagePublisher,
     var codeGenerationRemainingIterationCount: Int? = null,
     var codeGenerationTotalIterationCount: Int? = null,
-    var currentCodeGenerationId:  UUID? = null
-) : SessionState {
+    var currentCodeGenerationId:  UUID? = null,
+    override var token: CancellationTokenSource?
+    ) : SessionState {
     override val phase = SessionStatePhase.CODEGEN
 
     override suspend fun interact(action: SessionStateAction): SessionStateInteraction {
@@ -62,8 +62,7 @@ class CodeGenerationState(
                 tabId = tabID,
                 newPlaceholder = message("amazonqFeatureDev.code_generation.generating_code")
             )
-
-            val codeGenerationResult = generateCode(codeGenerationId = response.codeGenerationId(), messenger = messenger)
+            val codeGenerationResult = generateCode(codeGenerationId = response.codeGenerationId(), messenger = messenger, token = action.token)
             numberOfReferencesGenerated = codeGenerationResult.references.size
             numberOfFilesGenerated = codeGenerationResult.newFiles.size
             codeGenerationRemainingIterationCount = codeGenerationResult.codeGenerationRemainingIterationCount
@@ -115,11 +114,16 @@ class CodeGenerationState(
     }
 }
 
-private suspend fun CodeGenerationState.generateCode(codeGenerationId: String, messenger: MessagePublisher): CodeGenerationResult {
+private suspend fun CodeGenerationState.generateCode(codeGenerationId: String, messenger: MessagePublisher, token: CancellationTokenSource?): CodeGenerationResult {
     val pollCount = 180
     val requestDelay = 10000L
 
+
     repeat(pollCount) {
+
+        if (token?.token()?.isCancellationRequested == true) {
+            return CodeGenerationResult(emptyList(), emptyList(), emptyList());
+        }
         val codeGenerationResultState = config.featureDevService.getTaskAssistCodeGeneration(
             conversationId = config.conversationId,
             codeGenerationId = codeGenerationId,

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

@@ -8,7 +8,7 @@ import org.gradle.tooling.CancellationTokenSource
 class ConversationNotStartedState(
     override var approach: String,
     override val tabID: String,
-    override val token: CancellationTokenSource?,
+    override var token: CancellationTokenSource?,
     ) : SessionState {
     override val phase = SessionStatePhase.INIT
 

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

@@ -23,7 +23,7 @@ private val logger = getLogger<PrepareCodeGenerationState>()
 
 class PrepareCodeGenerationState(
     override var tabID: String,
-    override val token: CancellationTokenSource?,
+    override var token: CancellationTokenSource?,
     override var approach: String,
     private var config: SessionStateConfig,
     val filePaths: List<NewFileZipInfo>,

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

@@ -6,6 +6,7 @@ package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VfsUtil
+import org.gradle.tooling.GradleConnector
 import software.aws.toolkits.jetbrains.services.amazonq.FeatureDevSessionContext
 import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CODE_GENERATION_RETRY_LIMIT
@@ -81,7 +82,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
+            token = GradleConnector.newCancellationTokenSource()
         )
     }
 
@@ -115,6 +116,7 @@ class Session(val tabID: String, val project: Project) {
         var action = SessionStateAction(
             task = task,
             msg = msg,
+            token = sessionState.token
         )
         val resp = sessionState.interact(action)
         if (resp.nextState != null) {

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

@@ -8,7 +8,7 @@ import org.gradle.tooling.CancellationTokenSource
 interface SessionState {
     val tabID: String
     val phase: SessionStatePhase?
-    val token: CancellationTokenSource?
+    var token: CancellationTokenSource?
     var approach: String
     suspend fun interact(action: SessionStateAction): SessionStateInteraction
 }

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

@@ -4,13 +4,15 @@
 package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session
 
 import com.fasterxml.jackson.annotation.JsonValue
+import org.gradle.tooling.CancellationTokenSource
 import software.aws.toolkits.jetbrains.services.amazonq.FeatureDevSessionContext
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FeatureDevService
 import software.aws.toolkits.jetbrains.services.cwc.messages.RecommendationContentSpan
 
 data class SessionStateAction(
     val task: String,
     val msg: String,
+    val token: CancellationTokenSource? = null
 )
 
 data class Interaction(

plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationStateTest.kt

@@ -55,7 +55,8 @@ class CodeGenerationStateTest : FeatureDevTestBase() {
             testUploadId,
             0,
             testRepositorySize,
-            messenger
+            messenger,
+            token = null
         )
 
         mockkStatic(MessagePublisher::sendAnswerPart)

plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties

@@ -55,6 +55,8 @@ amazonqFeatureDev.code_generation.notification_message=Your code suggestions fro
 amazonqFeatureDev.code_generation.notification_open_link=Open chat
 amazonqFeatureDev.code_generation.notification_title=Amazon Q Developer Agent for software development
 amazonqFeatureDev.code_generation.provide_code_feedback=How can I improve the code for your use case?
+amazonqFeatureDev.code_generation.stopping_code_generation=Stopping code generation ...
+amazonqFeatureDev.code_generation.stopped_code_generation=You've stopped code generation.
 amazonqFeatureDev.code_generation.updated_code=Okay, I updated your code files. Would you like to work on another task?
 amazonqFeatureDev.content_length.error_text=The folder you selected is too large for me to use as context. Please choose a smaller folder to work on. For more information on quotas, see the [Amazon Q Developer documentation](https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/software-dev.html#quotas). 
 amazonqFeatureDev.error_text=Sorry, we encountered a problem when processing your request.