Fix

Author: siakmun-aws

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

@@ -55,22 +55,29 @@ export abstract class CodeGenBase {
         this.currentCodeGenerationId = config.currentCodeGenerationId || EmptyCodeGenID
     }
 
-    protected abstract handleProgress(messenger: BaseMessenger, detail?: string): void
+    protected abstract handleProgress(messenger: BaseMessenger, action: SessionStateAction, detail?: string): void
     protected abstract getScheme(): string
-    protected abstract handleGenerationComplete(messenger: BaseMessenger, newFileInfo: NewFileInfo[]): void
+    protected abstract getTimeoutErrorCode(): string
+    protected abstract handleGenerationComplete(
+        messenger: BaseMessenger,
+        newFileInfo: NewFileInfo[],
+        action: SessionStateAction
+    ): void
 
     async generateCode({
         messenger,
         fs,
         codeGenerationId,
         telemetry: telemetry,
         workspaceFolders,
+        action,
     }: {
         messenger: BaseMessenger
         fs: VirtualFileSystem
         codeGenerationId: string
         telemetry: any
         workspaceFolders: CurrentWsFolders
+        action: SessionStateAction
     }): Promise<{
         newFiles: NewFileInfo[]
         deletedFiles: DeletedFileInfo[]
@@ -106,7 +113,7 @@ export abstract class CodeGenBase {
                     )
                     telemetry.setNumberOfFilesGenerated(newFileInfo.length)
 
-                    this.handleGenerationComplete(messenger, newFileInfo)
+                    this.handleGenerationComplete(messenger, newFileInfo, action)
 
                     return {
                         newFiles: newFileInfo,
@@ -119,7 +126,7 @@ export abstract class CodeGenBase {
                 case CodeGenerationStatus.PREDICT_READY:
                 case CodeGenerationStatus.IN_PROGRESS: {
                     if (codegenResult.codeGenerationStatusDetail) {
-                        this.handleProgress(messenger, codegenResult.codeGenerationStatusDetail)
+                        this.handleProgress(messenger, action, codegenResult.codeGenerationStatusDetail)
                     }
                     await new Promise((f) => globals.clock.setTimeout(f, this.requestDelay))
                     break
@@ -138,7 +145,7 @@ export abstract class CodeGenBase {
 
         if (!this.isCancellationRequested) {
             const errorMessage = i18n('AWS.amazonq.featureDev.error.codeGen.timeout')
-            throw new ToolkitError(errorMessage, { code: 'CodeGenTimeout' })
+            throw new ToolkitError(errorMessage, { code: this.getTimeoutErrorCode() })
         }
 
         return {
@@ -210,7 +217,11 @@ export abstract class BasePrepareCodeGenState implements SessionState {
         )
     }
 
+    protected abstract preUpload(action: SessionStateAction): void
+    protected abstract postUpload(action: SessionStateAction): void
+
     async interact(action: SessionStateAction): Promise<SessionStateInteraction> {
+        this.preUpload(action)
         const uploadId = await telemetry.amazonq_createUpload.run(async (span) => {
             span.record({
                 amazonqConversationId: this.config.conversationId,
@@ -231,6 +242,7 @@ export abstract class BasePrepareCodeGenState implements SessionState {
             )
 
             await uploadCode(uploadUrl, zipFileBuffer, zipFileChecksum, kmsKeyArn)
+            this.postUpload(action)
 
             return uploadId
         })
@@ -338,6 +350,7 @@ export abstract class BaseCodeGenState extends CodeGenBase implements SessionSta
                     codeGenerationId,
                     telemetry: action.telemetry,
                     workspaceFolders: this.config.workspaceFolders,
+                    action,
                 })
 
                 if (codeGeneration && !action.tokenSource?.token.isCancellationRequested) {

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

@@ -25,14 +25,14 @@ import { BaseCodeGenState, BasePrepareCodeGenState, CreateNextStateParams } from
 import { Intent } from '../../amazonq/commons/types'
 
 export class DocCodeGenState extends BaseCodeGenState {
-    protected handleProgress(messenger: DocMessenger, detail?: string): void {
+    protected handleProgress(messenger: DocMessenger, action: SessionStateAction, detail?: string): void {
         if (detail) {
             const progress = getFileSummaryPercentage(detail)
             messenger.sendDocProgress(
                 this.tabID,
                 progress === 100 ? DocGenerationStep.GENERATING_ARTIFACTS : DocGenerationStep.SUMMARIZING_FILES,
                 progress,
-                (messenger as any).mode
+                action.mode
             )
         }
     }
@@ -41,8 +41,16 @@ export class DocCodeGenState extends BaseCodeGenState {
         return docScheme
     }
 
-    protected handleGenerationComplete(messenger: DocMessenger, newFileInfo: NewFileInfo[]): void {
-        messenger.sendDocProgress(this.tabID, DocGenerationStep.GENERATING_ARTIFACTS + 1, 100, (messenger as any).mode)
+    protected getTimeoutErrorCode(): string {
+        return 'DocGenerationTimeout'
+    }
+
+    protected handleGenerationComplete(
+        messenger: DocMessenger,
+        newFileInfo: NewFileInfo[],
+        action: SessionStateAction
+    ): void {
+        messenger.sendDocProgress(this.tabID, DocGenerationStep.GENERATING_ARTIFACTS + 1, 100, action.mode)
     }
 
     protected handleError(messenger: DocMessenger, codegenResult: any): Error {
@@ -118,6 +126,14 @@ export class DocCodeGenState extends BaseCodeGenState {
 }
 
 export class DocPrepareCodeGenState extends BasePrepareCodeGenState {
+    protected preUpload(action: SessionStateAction): void {
+        // Do nothing
+    }
+
+    protected postUpload(action: SessionStateAction): void {
+        // Do nothing
+    }
+
     protected override createNextState(config: SessionStateConfig): SessionState {
         return super.createNextState(config, DocCodeGenState)
     }

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

@@ -147,7 +147,7 @@ export class MockCodeGenState implements SessionState {
 }
 
 export class FeatureDevCodeGenState extends BaseCodeGenState {
-    protected handleProgress(messenger: Messenger, detail?: string): void {
+    protected handleProgress(messenger: Messenger, action: SessionStateAction, detail?: string): void {
         if (detail) {
             messenger.sendAnswer({
                 message: i18n('AWS.amazonq.featureDev.pillText.generatingCode') + `\n\n${detail}`,
@@ -161,7 +161,15 @@ export class FeatureDevCodeGenState extends BaseCodeGenState {
         return featureDevScheme
     }
 
-    protected handleGenerationComplete(_messenger: Messenger, _newFileInfo: NewFileInfo[]): void {
+    protected getTimeoutErrorCode(): string {
+        return 'CodeGenTimeout'
+    }
+
+    protected handleGenerationComplete(
+        _messenger: Messenger,
+        _newFileInfo: NewFileInfo[],
+        action: SessionStateAction
+    ): void {
         // No special handling needed for feature dev
     }
 
@@ -220,11 +228,40 @@ export class FeatureDevCodeGenState extends BaseCodeGenState {
     }
 
     protected override createNextState(config: SessionStateConfig, params: CreateNextStateParams): SessionState {
-        return super.createNextState(config, params, FeatureDevPrepareCodeGenState)
+        return super.createNextState(
+            { ...config, currentCodeGenerationId: this.currentCodeGenerationId },
+            params,
+            FeatureDevPrepareCodeGenState
+        )
     }
 }
 
 export class FeatureDevPrepareCodeGenState extends BasePrepareCodeGenState {
+    protected preUpload(action: SessionStateAction): void {
+        action.messenger.sendAnswer({
+            message: i18n('AWS.amazonq.featureDev.pillText.uploadingCode'),
+            type: 'answer-part',
+            tabID: this.tabID,
+        })
+
+        action.messenger.sendUpdatePlaceholder(this.tabID, i18n('AWS.amazonq.featureDev.pillText.uploadingCode'))
+    }
+
+    protected postUpload(action: SessionStateAction): void {
+        if (!action.tokenSource?.token.isCancellationRequested) {
+            action.messenger.sendAnswer({
+                message: i18n('AWS.amazonq.featureDev.pillText.contextGatheringCompleted'),
+                type: 'answer-part',
+                tabID: this.tabID,
+            })
+
+            action.messenger.sendUpdatePlaceholder(
+                this.tabID,
+                i18n('AWS.amazonq.featureDev.pillText.contextGatheringCompleted')
+            )
+        }
+    }
+
     protected override createNextState(config: SessionStateConfig): SessionState {
         return super.createNextState(config, FeatureDevCodeGenState)
     }