fix(dev): use currentIteration since the counts isnt provided when stop code gen

Author: tverney

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

@@ -57,7 +57,7 @@ suspend fun FeatureDevController.onCodeGeneration(
         var totalIterations: Int? = state.codeGenerationTotalIterationCount
 
         if (state.token?.token()?.isCancellationRequested == true) {
-            this.disposeToken(state, messenger, tabId, remainingIterations, totalIterations)
+            this.disposeToken(state, messenger, tabId, state.currentIteration?.let { CODE_GENERATION_RETRY_LIMIT.minus(it) }, CODE_GENERATION_RETRY_LIMIT)
             return
         }
 
@@ -82,7 +82,7 @@ suspend fun FeatureDevController.onCodeGeneration(
         }
 
         if (state.token?.token()?.isCancellationRequested == true) {
-            disposeToken(state, messenger, tabId, remainingIterations, totalIterations)
+            disposeToken(state, messenger, tabId, state.currentIteration?.let { CODE_GENERATION_RETRY_LIMIT.minus(it) }, CODE_GENERATION_RETRY_LIMIT)
             return
         }
 
@@ -162,25 +162,50 @@ private suspend fun FeatureDevController.disposeToken(
     remainingIterations: Number?,
     totalIterations: Number?,
 ) {
-    if (state.codeGenerationRemainingIterationCount !== null) {
+
+    if (remainingIterations !== null &&  remainingIterations.toInt() <= 0) {
         messenger.sendAnswer(
             tabId = tabId,
             messageType = FeatureDevMessageType.Answer,
             message =
             message(
-                "amazonqFeatureDev.code_generation.stopped_code_generation",
-                remainingIterations ?: state.currentIteration as Any,
-                totalIterations ?: CODE_GENERATION_RETRY_LIMIT,
+                "amazonqFeatureDev.code_generation.stopped_code_generation_no_iterations"
             ),
         )
-    } else {
-        messenger.sendAnswer(
+        //I stopped generating your code. You don't have more iterations left, however, you can start a new session
+        messenger.sendSystemPrompt(
             tabId = tabId,
-            messageType = FeatureDevMessageType.Answer,
-            message = message("amazonqFeatureDev.code_generation.stopped_code_generation_without_total", state.currentIteration as Any),
+            followUp = listOf(
+                FollowUp(
+                    pillText = message("amazonqFeatureDev.follow_up.new_task"),
+                    type = FollowUpTypes.NEW_TASK,
+                    status = FollowUpStatusType.Info
+                ),
+                FollowUp(
+                    pillText = message("amazonqFeatureDev.follow_up.close_session"),
+                    type = FollowUpTypes.CLOSE_SESSION,
+                    status = FollowUpStatusType.Info
+                )
+            )
         )
+        messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false)
+        messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.after_code_generation"))
+
+        return
     }
 
+
+    messenger.sendAnswer(
+        tabId = tabId,
+        messageType = FeatureDevMessageType.Answer,
+        message =
+        message(
+            "amazonqFeatureDev.code_generation.stopped_code_generation",
+            remainingIterations ?: state.currentIteration?.let { CODE_GENERATION_RETRY_LIMIT.minus(it) } as Any,
+            totalIterations ?: CODE_GENERATION_RETRY_LIMIT,
+        ),
+    )
+
     messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = true)
 
     messenger.sendUpdatePlaceholder(

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

@@ -58,8 +58,8 @@ 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.stopped_code_generation=I stopped generating your code. If you want to continue working on this task, provide another description. You have {0} out of {1} code generations left.
-amazonqFeatureDev.code_generation.stopped_code_generation_without_total=I stopped generating your code. If you want to continue working on this task, provide another description. You have started {0} code generations.
 amazonqFeatureDev.code_generation.stopping_code_generation=Stopping code generation...
+amazonqFeatureDev.code_generation.stopped_code_generation_no_iterations=I stopped generating your code. You don't have more iterations left, however, you can start a new session.
 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.