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

packages/amazonq/.changes/next-release/Feature-f03a4f2f-5819-4276-b914-be0b1a4ecedf.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Disable displaying remaining code generation iterations until 2 or less 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>
@@ -465,7 +463,7 @@ export class FeatureDevController {
                     type: 'answer',
                     tabID: tabID,
                     message:
-                        remainingIterations === 0
+                        remainingIterations > 2 || 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.`,
                 })
@@ -518,9 +516,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 +560,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,18 @@ 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 =
+                        remainingIterations > 2 || 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.`
+                    await verifyAddCodeMessage(remainingIterations, totalIterations, expectedMessage)
+                })
+            }
         })
 
         describe('processErrorChatMessage', function () {

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

@@ -111,17 +111,17 @@ describe('sessionState', () => {
         it('transitions to PrepareCodeGenState when codeGenerationStatus ready ', async () => {
             mockGetCodeGeneration = sinon.stub().resolves({
                 codeGenerationStatus: { status: 'Complete' },
-                codeGenerationRemainingIterationCount: 2,
-                codeGenerationTotalIterationCount: 3,
+                codeGenerationRemainingIterationCount: 9,
+                codeGenerationTotalIterationCount: 10,
             })
 
             mockExportResultArchive = sinon.stub().resolves({ newFileContents: [], deletedFiles: [], references: [] })
 
             const testAction = mockSessionStateAction()
-            const state = new CodeGenState(testConfig, [], [], [], tabId, 0, {}, 2, 3)
+            const state = new CodeGenState(testConfig, [], [], [], tabId, 0, {}, 9, 10)
             const result = await state.interact(testAction)
 
-            const nextState = new PrepareCodeGenState(testConfig, [], [], [], tabId, 1, 2, 3, undefined)
+            const nextState = new PrepareCodeGenState(testConfig, [], [], [], tabId, 1, 9, 10, undefined)
 
             assert.deepStrictEqual(result.nextState?.deletedFiles, nextState.deletedFiles)
             assert.deepStrictEqual(result.nextState?.filePaths, result.nextState?.filePaths)