fix(amazonq): use server-side data to decide remaining code generations #6370

## Problem
Previously remaining code generation iterations are always displayed.

## Solution
Only display when it goes down to 2 or less.

Author: siakmun-aws

packages/amazonq/.changes/next-release/Bug Fix-ef93f909-3aa5-4e62-a4fc-850376161d24.json

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

packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts

@@ -49,8 +49,6 @@ import { FollowUpTypes } from '../../../amazonq/commons/types'
 import { Messenger } from '../../../amazonq/commons/connector/baseMessenger'
 import { BaseChatSessionStorage } from '../../../amazonq/commons/baseChatStorage'
 
-export const TotalSteps = 3
-
 export interface ChatControllerEventEmitters {
     readonly processHumanChatMessage: EventEmitter<any>
     readonly followUpClicked: EventEmitter<any>
@@ -462,12 +460,17 @@ export class FeatureDevController {
 
             if (remainingIterations !== undefined && totalIterations !== undefined) {
                 this.messenger.sendAnswer({
-                    type: 'answer',
+                    type: 'answer' as const,
                     tabID: tabID,
-                    message:
-                        remainingIterations === 0
-                            ? 'Would you like me to add this code to your project?'
-                            : `Would you like me to add this code to your project, or provide feedback for new code? You have ${remainingIterations} out of ${totalIterations} code generations left.`,
+                    message: (() => {
+                        if (remainingIterations > 2) {
+                            return 'Would you like me to add this code to your project, or provide feedback for new code?'
+                        } else if (remainingIterations > 0) {
+                            return `Would you like me to add this code to your project, or provide feedback for new code? You have ${remainingIterations} out of ${totalIterations} code generations left.`
+                        } else {
+                            return 'Would you like me to add this code to your project?'
+                        }
+                    })(),
                 })
             }
 
@@ -518,9 +521,8 @@ export class FeatureDevController {
             if (session?.state?.tokenSource?.token.isCancellationRequested) {
                 this.workOnNewTask(
                     session.tabID,
-                    session.state.codeGenerationRemainingIterationCount ||
-                        TotalSteps - (session.state?.currentIteration || 0),
-                    session.state.codeGenerationTotalIterationCount || TotalSteps,
+                    session.state.codeGenerationRemainingIterationCount,
+                    session.state.codeGenerationTotalIterationCount,
                     session?.state?.tokenSource?.token.isCancellationRequested
                 )
                 this.disposeToken(session)
@@ -563,10 +565,16 @@ export class FeatureDevController {
     ) {
         if (isStoppedGeneration) {
             this.messenger.sendAnswer({
-                message:
-                    (remainingIterations ?? 0) <= 0
-                        ? "I stopped generating your code. You don't have more iterations left, however, you can start a new session."
-                        : `I stopped generating your code. If you want to continue working on this task, provide another description. You have ${remainingIterations} out of ${totalIterations} code generations left.`,
+                message: ((remainingIterations) => {
+                    if (remainingIterations && totalIterations) {
+                        if (remainingIterations <= 0) {
+                            return "I stopped generating your code. You don't have more iterations left, however, you can start a new session."
+                        } else if (remainingIterations <= 2) {
+                            return `I stopped generating your code. If you want to continue working on this task, provide another description. You have ${remainingIterations} out of ${totalIterations} code generations left.`
+                        }
+                    }
+                    return 'I stopped generating your code. If you want to continue working on this task, provide another description.'
+                })(remainingIterations),
                 type: 'answer-part',
                 tabID,
             })

packages/core/src/amazonqFeatureDev/session/sessionState.ts

@@ -175,14 +175,16 @@ abstract class CodeGenBase {
         codeGenerationRemainingIterationCount?: number
         codeGenerationTotalIterationCount?: number
     }> {
+        let codeGenerationRemainingIterationCount = undefined
+        let codeGenerationTotalIterationCount = undefined
         for (
             let pollingIteration = 0;
             pollingIteration < this.pollCount && !this.isCancellationRequested;
             ++pollingIteration
         ) {
             const codegenResult = await this.config.proxyClient.getCodeGeneration(this.conversationId, codeGenerationId)
-            const codeGenerationRemainingIterationCount = codegenResult.codeGenerationRemainingIterationCount
-            const codeGenerationTotalIterationCount = codegenResult.codeGenerationTotalIterationCount
+            codeGenerationRemainingIterationCount = codegenResult.codeGenerationRemainingIterationCount
+            codeGenerationTotalIterationCount = codegenResult.codeGenerationTotalIterationCount
 
             getLogger().debug(`Codegen response: %O`, codegenResult)
             telemetry.setCodeGenerationResult(codegenResult.codeGenerationStatus.status)
@@ -272,6 +274,8 @@ abstract class CodeGenBase {
             newFiles: [],
             deletedFiles: [],
             references: [],
+            codeGenerationRemainingIterationCount: codeGenerationRemainingIterationCount,
+            codeGenerationTotalIterationCount: codeGenerationTotalIterationCount,
         }
     }
 }
@@ -345,8 +349,13 @@ export class CodeGenState extends CodeGenBase implements SessionState {
                 this.filePaths = codeGeneration.newFiles
                 this.deletedFiles = codeGeneration.deletedFiles
                 this.references = codeGeneration.references
+
                 this.codeGenerationRemainingIterationCount = codeGeneration.codeGenerationRemainingIterationCount
                 this.codeGenerationTotalIterationCount = codeGeneration.codeGenerationTotalIterationCount
+                this.currentIteration =
+                    this.codeGenerationRemainingIterationCount && this.codeGenerationTotalIterationCount
+                        ? this.codeGenerationTotalIterationCount - this.codeGenerationRemainingIterationCount
+                        : this.currentIteration + 1
 
                 if (action.uploadHistory && !action.uploadHistory[codeGenerationId] && codeGenerationId) {
                     action.uploadHistory[codeGenerationId] = {
@@ -366,7 +375,7 @@ export class CodeGenState extends CodeGenBase implements SessionState {
                     this.deletedFiles,
                     this.references,
                     this.tabID,
-                    this.currentIteration + 1,
+                    this.currentIteration,
                     this.codeGenerationRemainingIterationCount,
                     this.codeGenerationTotalIterationCount,
                     action.uploadHistory,

packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts

@@ -526,6 +526,29 @@ describe('Controller', () => {
                 await waitUntil(() => Promise.resolve(sendMetricDataTelemetrySpy.callCount >= 2), {})
             }
 
+            async function verifyAddCodeMessage(
+                remainingIterations: number,
+                totalIterations: number,
+                expectedMessage: string
+            ) {
+                sinon.stub(session, 'send').resolves()
+                sinon.stub(session, 'sendLinesOfCodeGeneratedTelemetry').resolves() // Avoid sending extra telemetry
+                const sendAnswerSpy = sinon.stub(controllerSetup.messenger, 'sendAnswer')
+                sinon.stub(session.state, 'codeGenerationRemainingIterationCount').value(remainingIterations)
+                sinon.stub(session.state, 'codeGenerationTotalIterationCount').value(totalIterations)
+
+                await fireChatMessage(session)
+                await verifyMetricsCalled()
+
+                assert.ok(
+                    sendAnswerSpy.calledWith({
+                        type: 'answer',
+                        tabID,
+                        message: expectedMessage,
+                    })
+                )
+            }
+
             beforeEach(async () => {
                 session = await createCodeGenState()
                 sinon.stub(session, 'preloader').resolves()
@@ -558,6 +581,23 @@ describe('Controller', () => {
                     await verifyException(error)
                 })
             }
+
+            // Using 3 to avoid spamming the tests
+            for (let remainingIterations = 0; remainingIterations <= 3; remainingIterations++) {
+                it(`verifies add code messages for remaining iterations at ${remainingIterations}`, async () => {
+                    const totalIterations = 10
+                    const expectedMessage = (() => {
+                        if (remainingIterations > 2) {
+                            return 'Would you like me to add this code to your project, or provide feedback for new code?'
+                        } else if (remainingIterations > 0) {
+                            return `Would you like me to add this code to your project, or provide feedback for new code? You have ${remainingIterations} out of ${totalIterations} code generations left.`
+                        } else {
+                            return 'Would you like me to add this code to your project?'
+                        }
+                    })()
+                    await verifyAddCodeMessage(remainingIterations, totalIterations, expectedMessage)
+                })
+            }
         })
 
         describe('processErrorChatMessage', function () {