Skip to content

Commit 420d9da

Browse files
siakmun-awssiakmun-aws
committed
fix(amazonq): conditionally display remaining code generation iterations
1 parent c9b4746 commit 420d9da

File tree

6 files changed

+106
-62
lines changed

6 files changed

+106
-62
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "Feature",
3+
"description": "Disable displaying remaining code generation iterations until 2 or less remaining"
4+
}

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

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,8 +49,6 @@ import { FollowUpTypes } from '../../../amazonq/commons/types'
4949
import { Messenger } from '../../../amazonq/commons/connector/baseMessenger'
5050
import { BaseChatSessionStorage } from '../../../amazonq/commons/baseChatStorage'
5151

52-
export const TotalSteps = 3
53-
5452
export interface ChatControllerEventEmitters {
5553
readonly processHumanChatMessage: EventEmitter<any>
5654
readonly followUpClicked: EventEmitter<any>
@@ -465,7 +463,7 @@ export class FeatureDevController {
465463
type: 'answer',
466464
tabID: tabID,
467465
message:
468-
remainingIterations === 0
466+
remainingIterations > 2 || remainingIterations <= 0
469467
? 'Would you like me to add this code to your project?'
470468
: `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.`,
471469
})
@@ -518,9 +516,8 @@ export class FeatureDevController {
518516
if (session?.state?.tokenSource?.token.isCancellationRequested) {
519517
this.workOnNewTask(
520518
session.tabID,
521-
session.state.codeGenerationRemainingIterationCount ||
522-
TotalSteps - (session.state?.currentIteration || 0),
523-
session.state.codeGenerationTotalIterationCount || TotalSteps,
519+
session.state.codeGenerationRemainingIterationCount,
520+
session.state.codeGenerationTotalIterationCount,
524521
session?.state?.tokenSource?.token.isCancellationRequested
525522
)
526523
this.disposeToken(session)
@@ -563,10 +560,16 @@ export class FeatureDevController {
563560
) {
564561
if (isStoppedGeneration) {
565562
this.messenger.sendAnswer({
566-
message:
567-
(remainingIterations ?? 0) <= 0
568-
? "I stopped generating your code. You don't have more iterations left, however, you can start a new session."
569-
: `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.`,
563+
message: ((remainingIterations) => {
564+
if (remainingIterations && totalIterations) {
565+
if (remainingIterations <= 0) {
566+
return "I stopped generating your code. You don't have more iterations left, however, you can start a new session."
567+
} else if (remainingIterations <= 2) {
568+
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.`
569+
}
570+
}
571+
return 'I stopped generating your code. If you want to continue working on this task, provide another description.'
572+
})(remainingIterations),
570573
type: 'answer-part',
571574
tabID,
572575
})

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

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -175,14 +175,16 @@ abstract class CodeGenBase {
175175
codeGenerationRemainingIterationCount?: number
176176
codeGenerationTotalIterationCount?: number
177177
}> {
178+
let codeGenerationRemainingIterationCount = undefined
179+
let codeGenerationTotalIterationCount = undefined
178180
for (
179181
let pollingIteration = 0;
180182
pollingIteration < this.pollCount && !this.isCancellationRequested;
181183
++pollingIteration
182184
) {
183185
const codegenResult = await this.config.proxyClient.getCodeGeneration(this.conversationId, codeGenerationId)
184-
const codeGenerationRemainingIterationCount = codegenResult.codeGenerationRemainingIterationCount
185-
const codeGenerationTotalIterationCount = codegenResult.codeGenerationTotalIterationCount
186+
codeGenerationRemainingIterationCount = codegenResult.codeGenerationRemainingIterationCount
187+
codeGenerationTotalIterationCount = codegenResult.codeGenerationTotalIterationCount
186188

187189
getLogger().debug(`Codegen response: %O`, codegenResult)
188190
telemetry.setCodeGenerationResult(codegenResult.codeGenerationStatus.status)
@@ -272,6 +274,8 @@ abstract class CodeGenBase {
272274
newFiles: [],
273275
deletedFiles: [],
274276
references: [],
277+
codeGenerationRemainingIterationCount: codeGenerationRemainingIterationCount,
278+
codeGenerationTotalIterationCount: codeGenerationTotalIterationCount,
275279
}
276280
}
277281
}
@@ -345,8 +349,13 @@ export class CodeGenState extends CodeGenBase implements SessionState {
345349
this.filePaths = codeGeneration.newFiles
346350
this.deletedFiles = codeGeneration.deletedFiles
347351
this.references = codeGeneration.references
352+
348353
this.codeGenerationRemainingIterationCount = codeGeneration.codeGenerationRemainingIterationCount
349354
this.codeGenerationTotalIterationCount = codeGeneration.codeGenerationTotalIterationCount
355+
this.currentIteration =
356+
this.codeGenerationRemainingIterationCount && this.codeGenerationTotalIterationCount
357+
? this.codeGenerationTotalIterationCount - this.codeGenerationRemainingIterationCount
358+
: this.currentIteration + 1
350359

351360
if (action.uploadHistory && !action.uploadHistory[codeGenerationId] && codeGenerationId) {
352361
action.uploadHistory[codeGenerationId] = {
@@ -366,7 +375,7 @@ export class CodeGenState extends CodeGenBase implements SessionState {
366375
this.deletedFiles,
367376
this.references,
368377
this.tabID,
369-
this.currentIteration + 1,
378+
this.currentIteration,
370379
this.codeGenerationRemainingIterationCount,
371380
this.codeGenerationTotalIterationCount,
372381
action.uploadHistory,

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -526,6 +526,29 @@ describe('Controller', () => {
526526
await waitUntil(() => Promise.resolve(sendMetricDataTelemetrySpy.callCount >= 2), {})
527527
}
528528

529+
async function verifyAddCodeMessage(
530+
remainingIterations: number,
531+
totalIterations: number,
532+
expectedMessage: string
533+
) {
534+
sinon.stub(session, 'send').resolves()
535+
sinon.stub(session, 'sendLinesOfCodeGeneratedTelemetry').resolves() // Avoid sending extra telemetry
536+
const sendAnswerSpy = sinon.stub(controllerSetup.messenger, 'sendAnswer')
537+
sinon.stub(session.state, 'codeGenerationRemainingIterationCount').value(remainingIterations)
538+
sinon.stub(session.state, 'codeGenerationTotalIterationCount').value(totalIterations)
539+
540+
await fireChatMessage(session)
541+
await verifyMetricsCalled()
542+
543+
assert.ok(
544+
sendAnswerSpy.calledWith({
545+
type: 'answer',
546+
tabID,
547+
message: expectedMessage,
548+
})
549+
)
550+
}
551+
529552
beforeEach(async () => {
530553
session = await createCodeGenState()
531554
sinon.stub(session, 'preloader').resolves()
@@ -558,6 +581,18 @@ describe('Controller', () => {
558581
await verifyException(error)
559582
})
560583
}
584+
585+
// Using 3 to avoid spamming the tests
586+
for (let remainingIterations = 0; remainingIterations <= 3; remainingIterations++) {
587+
it(`verifies add code messages for remaining iterations at ${remainingIterations}`, async () => {
588+
const totalIterations = 10
589+
const expectedMessage =
590+
remainingIterations > 2 || remainingIterations <= 0
591+
? 'Would you like me to add this code to your project?'
592+
: `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.`
593+
await verifyAddCodeMessage(remainingIterations, totalIterations, expectedMessage)
594+
})
595+
}
561596
})
562597

563598
describe('processErrorChatMessage', function () {

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -111,17 +111,17 @@ describe('sessionState', () => {
111111
it('transitions to PrepareCodeGenState when codeGenerationStatus ready ', async () => {
112112
mockGetCodeGeneration = sinon.stub().resolves({
113113
codeGenerationStatus: { status: 'Complete' },
114-
codeGenerationRemainingIterationCount: 2,
115-
codeGenerationTotalIterationCount: 3,
114+
codeGenerationRemainingIterationCount: 9,
115+
codeGenerationTotalIterationCount: 10,
116116
})
117117

118118
mockExportResultArchive = sinon.stub().resolves({ newFileContents: [], deletedFiles: [], references: [] })
119119

120120
const testAction = mockSessionStateAction()
121-
const state = new CodeGenState(testConfig, [], [], [], tabId, 0, {}, 2, 3)
121+
const state = new CodeGenState(testConfig, [], [], [], tabId, 0, {}, 9, 10)
122122
const result = await state.interact(testAction)
123123

124-
const nextState = new PrepareCodeGenState(testConfig, [], [], [], tabId, 1, 2, 3, undefined)
124+
const nextState = new PrepareCodeGenState(testConfig, [], [], [], tabId, 1, 9, 10, undefined)
125125

126126
assert.deepStrictEqual(result.nextState?.deletedFiles, nextState.deletedFiles)
127127
assert.deepStrictEqual(result.nextState?.filePaths, result.nextState?.filePaths)

0 commit comments

Comments
 (0)