feat: proactively show code generation iterations

aws · Jul 10, 2024 · 5616b97 · 5616b97
1 parent 10dacd9
commit 5616b97
Show file tree

Hide file tree

Showing 7 changed files with 53 additions and 7 deletions.
diff --git a/packages/amazonq/.changes/next-release/Feature-4a1c951a-2289-4808-b106-86fe4f03528e.json b/packages/amazonq/.changes/next-release/Feature-4a1c951a-2289-4808-b106-86fe4f03528e.json
@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Amazon Q/dev: proactively show code generation iterations"
+}
diff --git a/packages/core/src/amazonqFeatureDev/client/codewhispererruntime-2022-11-11.json b/packages/core/src/amazonqFeatureDev/client/codewhispererruntime-2022-11-11.json
@@ -1262,6 +1262,12 @@
                 },
                 "codeGenerationStatus": {
                     "shape": "CodeGenerationStatus"
+                },
+                "codeGenerationRemainingIterationCount": {
+                    "shape": "Integer"
+                },
+                "codeGenerationTotalIterationCount": {
+                    "shape": "Integer"
                 }
             },
             "documentation": "<p>Response for getting task assist code generation status.</p>"

diff --git a/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts b/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts
@@ -476,6 +476,18 @@ export class FeatureDevController {
                 tabID,
                 session.uploadId
             )
+
+            const remainingIterations = session.state.codeGenerationRemainingIterationCount
+            const totalIterations = session.state.codeGenerationTotalIterationCount
+
+            if (remainingIterations !== undefined && totalIterations !== undefined) {
+                this.messenger.sendAnswer({
+                    type: 'answer',
+                    tabID: tabID,
+                    message: `You have ${remainingIterations} out of ${totalIterations} code iterations remaining.`,
+                })
+            }
+
             this.messenger.sendAnswer({
                 message: undefined,
                 type: 'system-prompt',

diff --git a/packages/core/src/amazonqFeatureDev/session/sessionState.ts b/packages/core/src/amazonqFeatureDev/session/sessionState.ts
@@ -254,13 +254,18 @@ abstract class CodeGenBase {
         newFiles: NewFileInfo[]
         deletedFiles: DeletedFileInfo[]
         references: CodeReference[]
+        codeGenerationRemainingIterationCount?: number
+        codeGenerationTotalIterationCount?: number
     }> {
         for (
             let pollingIteration = 0;
             pollingIteration < this.pollCount && !this.tokenSource.token.isCancellationRequested;
             ++pollingIteration
         ) {
             const codegenResult = await this.config.proxyClient.getCodeGeneration(this.conversationId, codeGenerationId)
+            const codeGenerationRemainingIterationCount = codegenResult.codeGenerationRemainingIterationCount
+            const codeGenerationTotalIterationCount = codegenResult.codeGenerationTotalIterationCount
+
             getLogger().debug(`Codegen response: %O`, codegenResult)
             telemetry.setCodeGenerationResult(codegenResult.codeGenerationStatus.status)
             switch (codegenResult.codeGenerationStatus.status) {
@@ -273,6 +278,8 @@ abstract class CodeGenBase {
                         newFiles: newFileInfo,
                         deletedFiles: getDeletedFileInfos(deletedFiles, workspaceFolders),
                         references,
+                        codeGenerationRemainingIterationCount: codeGenerationRemainingIterationCount,
+                        codeGenerationTotalIterationCount: codeGenerationTotalIterationCount,
                     }
                 }
                 case 'predict-ready':
@@ -321,7 +328,9 @@ export class CodeGenState extends CodeGenBase implements SessionState {
         public deletedFiles: DeletedFileInfo[],
         public references: CodeReference[],
         tabID: string,
-        private currentIteration: number
+        private currentIteration: number,
+        public codeGenerationRemainingIterationCount?: number,
+        public codeGenerationTotalIterationCount?: number
     ) {
         super(config, tabID)
     }
@@ -358,6 +367,9 @@ 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
+
                 action.telemetry.setAmazonqNumberOfReferences(this.references.length)
                 action.telemetry.recordUserCodeGenerationTelemetry(span, this.conversationId)
                 const nextState = new PrepareCodeGenState(
@@ -367,7 +379,9 @@ export class CodeGenState extends CodeGenBase implements SessionState {
                     this.deletedFiles,
                     this.references,
                     this.tabID,
-                    this.currentIteration + 1
+                    this.currentIteration + 1,
+                    this.codeGenerationRemainingIterationCount,
+                    this.codeGenerationTotalIterationCount
                 )
                 return {
                     nextState,
@@ -480,7 +494,9 @@ export class PrepareCodeGenState implements SessionState {
         public deletedFiles: DeletedFileInfo[],
         public references: CodeReference[],
         public tabID: string,
-        private currentIteration: number
+        private currentIteration: number,
+        public codeGenerationRemainingIterationCount?: number,
+        public codeGenerationTotalIterationCount?: number
     ) {
         this.tokenSource = new vscode.CancellationTokenSource()
         this.uploadId = config.uploadId

diff --git a/packages/core/src/amazonqFeatureDev/types.ts b/packages/core/src/amazonqFeatureDev/types.ts
@@ -57,6 +57,8 @@ export interface SessionState {
     readonly tabID: string
     interact(action: SessionStateAction): Promise<SessionStateInteraction>
     updateWorkspaceRoot?: (workspaceRoot: string) => void
+    codeGenerationRemainingIterationCount?: number
+    codeGenerationTotalIterationCount?: number
 }
 
 export interface SessionStateConfig {

diff --git a/packages/core/src/codewhisperer/client/user-service-2.json b/packages/core/src/codewhisperer/client/user-service-2.json
@@ -1081,7 +1081,9 @@
             "members": {
                 "conversationId": { "shape": "ConversationId" },
                 "codeGenerationStatus": { "shape": "CodeGenerationStatus" },
-                "codeGenerationStatusDetail": { "shape": "CodeGenerationStatusDetail" }
+                "codeGenerationStatusDetail": { "shape": "CodeGenerationStatusDetail" },
+                "codeGenerationRemainingIterationCount": { "shape": "Integer" },
+                "codeGenerationTotalIterationCount": { "shape": "Integer" }
             },
             "documentation": "<p>Response for getting task assist code generation status.</p>"
         },

diff --git a/packages/core/src/test/amazonqFeatureDev/session/sessionState.test.ts b/packages/core/src/test/amazonqFeatureDev/session/sessionState.test.ts
@@ -203,14 +203,18 @@ describe('sessionState', () => {
 
     describe('CodeGenState', () => {
         it('transitions to PrepareCodeGenState when codeGenerationStatus ready ', async () => {
-            mockGetCodeGeneration = sinon.stub().resolves({ codeGenerationStatus: { status: 'Complete' } })
+            mockGetCodeGeneration = sinon.stub().resolves({
+                codeGenerationStatus: { status: 'Complete' },
+                codeGenerationRemainingIterationCount: 2,
+                codeGenerationTotalIterationCount: 3,
+            })
             mockExportResultArchive = sinon.stub().resolves({ newFileContents: [], deletedFiles: [], references: [] })
 
             const testAction = mockSessionStateAction()
-            const state = new CodeGenState(testConfig, testApproach, [], [], [], tabId, 0)
+            const state = new CodeGenState(testConfig, testApproach, [], [], [], tabId, 0, 2, 3)
             const result = await state.interact(testAction)
 
-            const nextState = new PrepareCodeGenState(testConfig, testApproach, [], [], [], tabId, 1)
+            const nextState = new PrepareCodeGenState(testConfig, testApproach, [], [], [], tabId, 1, 2, 3)
 
             assert.deepStrictEqual(result, {
                 nextState,