fix(amazonq): rely on server-side data to communicate number of code … (aws#5263)

* fix(amazonq): rely on server-side data to communicate number of code generations remaining

* fix(amazonq): rely on server-side data to communicate number of code generations remaining

* fix(amazonq): rely on server-side data to communicate number of code generations remaining

* fix(amazonq): rely on server-side data to communicate number of code generations remaining

* fix(amazonq): rely on server-side data to communicate number of code generations remaining

Author: siakmun-aws

.changes/next-release/bugfix-693e3f08-15f7-4453-b304-e473c58bf46b.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Amazon Q /dev: Remove hard-coded limits and instead rely server-side data to communicate number of code generations remaining"
+}

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

@@ -28,7 +28,6 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Delete
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.NewFileZipInfo
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.PrepareCodeGenerationState
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Session
-import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionState
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.InsertAction
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions
@@ -64,7 +63,7 @@ suspend fun FeatureDevController.onCodeGeneration(
         var totalIterations: Int? = state.codeGenerationTotalIterationCount
 
         if (state.token?.token?.isCancellationRequested() == true) {
-            disposeToken(state, messenger, tabId, state.currentIteration?.let { CODE_GENERATION_RETRY_LIMIT.minus(it) }, CODE_GENERATION_RETRY_LIMIT)
+            disposeToken(messenger, tabId, state.codeGenerationRemainingIterationCount, state.codeGenerationTotalIterationCount)
             return
         }
 
@@ -96,7 +95,7 @@ suspend fun FeatureDevController.onCodeGeneration(
         }
 
         if (state.token?.token?.isCancellationRequested() == true) {
-            disposeToken(state, messenger, tabId, state.currentIteration?.let { CODE_GENERATION_RETRY_LIMIT.minus(it) }, CODE_GENERATION_RETRY_LIMIT)
+            disposeToken(messenger, tabId, state.codeGenerationRemainingIterationCount, state.codeGenerationTotalIterationCount)
             return
         }
 
@@ -133,14 +132,20 @@ suspend fun FeatureDevController.onCodeGeneration(
                 tabId = tabId,
                 messageType = FeatureDevMessageType.Answer,
                 message =
-                if (remainingIterations == 0) {
-                    message("amazonqFeatureDev.code_generation.iteration_zero")
-                } else {
+                if (remainingIterations > 2) {
+                    message("amazonqFeatureDev.code_generation.iteration_counts_ask_to_add_code_or_feedback")
+                } else if (remainingIterations > 0) {
                     message(
                         "amazonqFeatureDev.code_generation.iteration_counts",
                         remainingIterations,
                         totalIterations,
                     )
+                } else {
+                    message(
+                        "amazonqFeatureDev.code_generation.iteration_counts_ask_to_add_code",
+                        remainingIterations,
+                        totalIterations,
+                    )
                 },
             )
         }
@@ -196,7 +201,6 @@ suspend fun FeatureDevController.onCodeGeneration(
 }
 
 private suspend fun disposeToken(
-    state: SessionState,
     messenger: MessagePublisher,
     tabId: String,
     remainingIterations: Int?,
@@ -234,16 +238,25 @@ private suspend fun disposeToken(
         return
     }
 
-    messenger.sendAnswer(
-        tabId = tabId,
-        messageType = FeatureDevMessageType.Answer,
-        message =
-        message(
-            "amazonqFeatureDev.code_generation.stopped_code_generation",
-            remainingIterations ?: state.currentIteration?.let { CODE_GENERATION_RETRY_LIMIT - it } as Any,
-            totalIterations ?: CODE_GENERATION_RETRY_LIMIT,
-        ),
-    )
+    if (remainingIterations !== null && totalIterations !== null && remainingIterations <= 2) {
+        messenger.sendAnswer(
+            tabId = tabId,
+            messageType = FeatureDevMessageType.Answer,
+            message =
+            message(
+                "amazonqFeatureDev.code_generation.stopped_code_generation",
+                remainingIterations,
+                totalIterations,
+            ),
+        )
+    } else {
+        messenger.sendAnswer(
+            tabId = tabId,
+            messageType = FeatureDevMessageType.Answer,
+            message =
+            message("amazonqFeatureDev.code_generation.stopped_code_generation_no_iteration_count_display"),
+        )
+    }
 
     messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = true)
 

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

@@ -84,6 +84,12 @@ class CodeGenerationState(
             numberOfFilesGenerated = codeGenerationResult.newFiles.size
             codeGenerationRemainingIterationCount = codeGenerationResult.codeGenerationRemainingIterationCount
             codeGenerationTotalIterationCount = codeGenerationResult.codeGenerationTotalIterationCount
+            currentIteration =
+                if (codeGenerationRemainingIterationCount != null && codeGenerationTotalIterationCount != null) {
+                    codeGenerationTotalIterationCount?.let { total -> codeGenerationRemainingIterationCount?.let { remaining -> total - remaining } }
+                } else {
+                    currentIteration?.plus(1)
+                }
 
             runCatching {
                 var insertedLines = 0
@@ -128,7 +134,7 @@ class CodeGenerationState(
                     filePaths = codeGenerationResult.newFiles,
                     deletedFiles = codeGenerationResult.deletedFiles,
                     references = codeGenerationResult.references,
-                    currentIteration = currentIteration?.plus(1),
+                    currentIteration = currentIteration,
                     uploadId = uploadId,
                     messenger = messenger,
                     codeGenerationRemainingIterationCount = codeGenerationRemainingIterationCount,
@@ -182,17 +188,22 @@ private suspend fun CodeGenerationState.generateCode(
 ): CodeGenerationResult {
     val pollCount = 360
     val requestDelay = 5000L
+    var codeGenerationRemainingIterationCount: Int? = null
+    var codeGenerationTotalIterationCount: Int? = null
 
     repeat(pollCount) {
         if (token?.token?.isCancellationRequested() == true) {
-            return CodeGenerationResult(emptyList(), emptyList(), emptyList())
+            return CodeGenerationResult(emptyList(), emptyList(), emptyList(), codeGenerationRemainingIterationCount, codeGenerationTotalIterationCount)
         }
         val codeGenerationResultState =
             config.featureDevService.getTaskAssistCodeGeneration(
                 conversationId = config.conversationId,
                 codeGenerationId = codeGenerationId,
             )
 
+        codeGenerationRemainingIterationCount = codeGenerationResultState.codeGenerationRemainingIterationCount()
+        codeGenerationTotalIterationCount = codeGenerationResultState.codeGenerationTotalIterationCount()
+
         when (codeGenerationResultState.codeGenerationStatus().status()) {
             CodeGenerationWorkflowStatus.COMPLETE -> {
                 val codeGenerationStreamResult =
@@ -207,8 +218,8 @@ private suspend fun CodeGenerationState.generateCode(
                     newFiles = newFileInfo,
                     deletedFiles = deletedFileInfo,
                     references = codeGenerationStreamResult.references,
-                    codeGenerationRemainingIterationCount = codeGenerationResultState.codeGenerationRemainingIterationCount(),
-                    codeGenerationTotalIterationCount = codeGenerationResultState.codeGenerationTotalIterationCount(),
+                    codeGenerationRemainingIterationCount = codeGenerationRemainingIterationCount,
+                    codeGenerationTotalIterationCount = codeGenerationTotalIterationCount,
                 )
             }
             CodeGenerationWorkflowStatus.IN_PROGRESS -> {
@@ -252,7 +263,7 @@ private suspend fun CodeGenerationState.generateCode(
         }
     }
 
-    return CodeGenerationResult(emptyList(), emptyList(), emptyList())
+    return CodeGenerationResult(emptyList(), emptyList(), emptyList(), codeGenerationRemainingIterationCount, codeGenerationTotalIterationCount)
 }
 
 fun registerNewFiles(newFileContents: Map<String, String>): List<NewFileZipInfo> =

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

@@ -9,8 +9,8 @@ class ConversationNotStartedState(
     override var approach: String,
     override val tabID: String,
     override var token: CancellationTokenSource?,
-    override var codeGenerationRemainingIterationCount: Int?,
-    override var codeGenerationTotalIterationCount: Int?,
+    override var codeGenerationRemainingIterationCount: Int? = null,
+    override var codeGenerationTotalIterationCount: Int? = null,
     override var currentIteration: Int?,
     override var diffMetricsProcessed: DiffMetricsProcessed,
 ) : SessionState {

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

@@ -55,8 +55,6 @@ class Session(val tabID: String, val project: Project) {
             approach = "",
             tabID = tabID,
             token = null,
-            codeGenerationRemainingIterationCount = 0,
-            codeGenerationTotalIterationCount = CODE_GENERATION_RETRY_LIMIT,
             currentIteration = 0,
             diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet())
         )

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

@@ -85,15 +85,17 @@ amazonqFeatureDev.code_generation.error_message=I'm sorry, I ran into an issue w
 amazonqFeatureDev.code_generation.failed_generation=Code generation failed
 amazonqFeatureDev.code_generation.generating_code=Generating code ...
 amazonqFeatureDev.code_generation.iteration_counts=Would you like me to add this code to your project, or provide feedback for new code? You have {0} out of {1} code generations left.
+amazonqFeatureDev.code_generation.iteration_counts_ask_to_add_code=Would you like me to add this code to your project?
+amazonqFeatureDev.code_generation.iteration_counts_ask_to_add_code_or_feedback=Would you like me to add this code to your project, or provide feedback for new code?
 amazonqFeatureDev.code_generation.iteration_limit.error_text=Sorry, you've reached the quota for number of iterations on code generation. You can insert this code in your files or discuss a new plan. 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.code_generation.iteration_zero=Would you like me to add this code to your project?
 amazonqFeatureDev.code_generation.no_file_changes=Unable to generate any file changes
 amazonqFeatureDev.code_generation.no_retries.error_message=I'm sorry, I'm having trouble generating your code and can't continue at the moment. Please try again later, and share feedback to help me improve.
 amazonqFeatureDev.code_generation.notification_message=Your code suggestions from Amazon Q are ready to review
 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_no_iteration_count_display=I stopped generating your code. If you want to continue working on this task, provide another description.
 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.stopping_code_generation=Stopping code generation...
 amazonqFeatureDev.code_generation.updated_code=Okay, I updated your code files. Would you like to work on another task?