diff --git a/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts b/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts index 9ad24069f70..aad70595366 100644 --- a/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts +++ b/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts @@ -566,7 +566,7 @@ export class FeatureDevController { if (isStoppedGeneration) { this.messenger.sendAnswer({ message: ((remainingIterations) => { - if (remainingIterations && totalIterations) { + if (totalIterations !== undefined) { 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) { diff --git a/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts b/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts index 97044d99338..a6d36942840 100644 --- a/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts +++ b/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts @@ -43,6 +43,7 @@ import { featureDevScheme, featureName, messageWithConversationId } from '../../ import { i18n } from '../../../../shared/i18n-helper' import { FollowUpTypes } from '../../../../amazonq/commons/types' import { ToolkitError } from '../../../../shared' +import { MessengerTypes } from '../../../../amazonqFeatureDev/controllers/chat/messenger/constants' let mockGetCodeGeneration: sinon.SinonStub describe('Controller', () => { @@ -526,10 +527,11 @@ describe('Controller', () => { await waitUntil(() => Promise.resolve(sendMetricDataTelemetrySpy.callCount >= 2), {}) } - async function verifyAddCodeMessage( - remainingIterations: number, - totalIterations: number, - expectedMessage: string + async function verifyMessage( + expectedMessage: string, + type: MessengerTypes, + remainingIterations?: number, + totalIterations?: number ) { sinon.stub(session, 'send').resolves() sinon.stub(session, 'sendLinesOfCodeGeneratedTelemetry').resolves() // Avoid sending extra telemetry @@ -542,7 +544,7 @@ describe('Controller', () => { assert.ok( sendAnswerSpy.calledWith({ - type: 'answer', + type, tabID, message: expectedMessage, }) @@ -595,7 +597,33 @@ describe('Controller', () => { return 'Would you like me to add this code to your project?' } })() - await verifyAddCodeMessage(remainingIterations, totalIterations, expectedMessage) + await verifyMessage(expectedMessage, 'answer', remainingIterations, totalIterations) + }) + } + + for (let remainingIterations = -1; remainingIterations <= 3; remainingIterations++) { + let remaining: number | undefined = remainingIterations + if (remainingIterations < 0) { + remaining = undefined + } + it(`verifies messages after cancellation for remaining iterations at ${remaining !== undefined ? remaining : 'undefined'}`, async () => { + const totalIterations = 10 + const expectedMessage = (() => { + if (remaining === undefined || remaining > 2) { + return 'I stopped generating your code. If you want to continue working on this task, provide another description.' + } else if (remaining > 0) { + return `I stopped generating your code. If you want to continue working on this task, provide another description. You have ${remaining} out of ${totalIterations} code generations left.` + } else { + return "I stopped generating your code. You don't have more iterations left, however, you can start a new session." + } + })() + session.state.tokenSource.cancel() + await verifyMessage( + expectedMessage, + 'answer-part', + remaining, + remaining === undefined ? undefined : totalIterations + ) }) } })