fix(amazonqFeatureDev): include currentCodeGenerationId to track reference for stop code generation

Author: tverney

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

@@ -349,13 +349,8 @@ export class FeatureDevController {
             await session.send(message)
             const filePaths = session.state.filePaths ?? []
             const deletedFiles = session.state.deletedFiles ?? []
-
             // Only add the follow up accept/deny buttons when the tab hasn't been closed/request hasn't been cancelled
             if (session?.state?.tokenSource?.token.isCancellationRequested) {
-                session?.state.tokenSource?.dispose()
-                if (session?.state?.tokenSource) {
-                    session.state.tokenSource = undefined
-                }
                 return
             }
 
@@ -416,6 +411,14 @@ export class FeatureDevController {
             this.messenger.sendUpdatePlaceholder(tabID, i18n('AWS.amazonq.featureDev.pillText.selectOption'))
         } finally {
             // Finish processing the event
+
+            if (session?.state?.tokenSource?.token.isCancellationRequested) {
+                session?.state.tokenSource?.dispose()
+                if (session?.state?.tokenSource) {
+                    session.state.tokenSource = undefined
+                }
+                return
+            }
             this.messenger.sendAsyncEventProgress(tabID, false, undefined)
 
             // Lock the chat input until they explicitly click one of the follow ups

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

@@ -89,6 +89,7 @@ export class Session {
                 ...this.getSessionStateConfig(),
                 conversationId: this.conversationId,
                 uploadId: '',
+                currentCodeGenerationId: undefined,
             },
             [],
             [],
@@ -109,7 +110,6 @@ export class Session {
             workspaceFolders: this.config.workspaceFolders,
             proxyClient: this.proxyClient,
             conversationId: this.conversationId,
-            currentCodeGenerationId: this.currentCodeGenerationId as string,
         }
     }
 

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

@@ -126,7 +126,7 @@ abstract class CodeGenBase {
     public phase: SessionStatePhase = DevPhase.CODEGEN
     public readonly conversationId: string
     public readonly uploadId: string
-    public readonly currentCodeGenerationId?: string
+    public currentCodeGenerationId?: string
 
     constructor(
         protected config: SessionStateConfig,
@@ -274,6 +274,7 @@ export class CodeGenState extends CodeGenBase implements SessionState {
 
                 action.tokenSource?.token.onCancellationRequested(() => {
                     this.isCancellationRequested = true
+                    if (action.tokenSource) this.tokenSource = action.tokenSource
                     action.tokenSource?.dispose()
                     action.tokenSource = undefined
                 })
@@ -285,9 +286,12 @@ export class CodeGenState extends CodeGenBase implements SessionState {
                     this.config.conversationId,
                     this.config.uploadId,
                     action.msg,
-                    this.config.currentCodeGenerationId
+                    this.currentCodeGenerationId
                 )
 
+                this.currentCodeGenerationId = codeGenerationId
+                this.config.currentCodeGenerationId = codeGenerationId
+
                 if (!this.isCancellationRequested) {
                     action.messenger.sendAnswer({
                         message: i18n('AWS.amazonq.featureDev.pillText.generatingCode'),
@@ -325,7 +329,8 @@ export class CodeGenState extends CodeGenBase implements SessionState {
                     this.currentIteration + 1,
                     this.codeGenerationRemainingIterationCount,
                     this.codeGenerationTotalIterationCount,
-                    this.tokenSource
+                    this.tokenSource,
+                    this.currentCodeGenerationId
                 )
                 return {
                     nextState,
@@ -428,7 +433,6 @@ export class MockCodeGenState implements SessionState {
 export class PrepareCodeGenState implements SessionState {
     public readonly phase = DevPhase.CODEGEN
     public uploadId: string
-    public currentCodeGenerationId?: string
     public conversationId: string
     public tokenSource: vscode.CancellationTokenSource
     constructor(
@@ -440,11 +444,12 @@ export class PrepareCodeGenState implements SessionState {
         private currentIteration: number,
         public codeGenerationRemainingIterationCount?: number,
         public codeGenerationTotalIterationCount?: number,
-        public superTokenSource?: vscode.CancellationTokenSource
+        public superTokenSource?: vscode.CancellationTokenSource,
+        public currentCodeGenerationId?: string
     ) {
         this.tokenSource = superTokenSource || new vscode.CancellationTokenSource()
         this.uploadId = config.uploadId
-        this.currentCodeGenerationId = config.currentCodeGenerationId
+        this.currentCodeGenerationId = currentCodeGenerationId
         this.conversationId = config.conversationId
     }
 
@@ -496,7 +501,7 @@ export class PrepareCodeGenState implements SessionState {
             this.uploadId = uploadId
         }
         const nextState = new CodeGenState(
-            { ...this.config, uploadId: this.uploadId, currentCodeGenerationId: this.config.currentCodeGenerationId },
+            { ...this.config, uploadId: this.uploadId, currentCodeGenerationId: this.currentCodeGenerationId },
             this.filePaths,
             this.deletedFiles,
             this.references,

packages/core/src/amazonqFeatureDev/types.ts

@@ -21,6 +21,7 @@ export type Interaction = {
 export interface SessionStateInteraction {
     nextState: SessionState | Omit<SessionState, 'uploadId'> | undefined
     interaction: Interaction
+    currentCodeGenerationId?: string
 }
 
 export enum DevPhase {
@@ -60,7 +61,7 @@ export interface SessionState {
     readonly references?: CodeReference[]
     readonly phase?: SessionStatePhase
     readonly uploadId: string
-    readonly currentCodeGenerationId?: string
+    currentCodeGenerationId?: string
     tokenSource?: CancellationTokenSource
     readonly tabID: string
     interact(action: SessionStateAction): Promise<SessionStateInteraction>

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

@@ -30,7 +30,6 @@ const mockSessionStateAction = (msg?: string): SessionStateAction => {
     }
 }
 
-let mockGeneratePlan: sinon.SinonStub
 let mockGetCodeGeneration: sinon.SinonStub
 let mockExportResultArchive: sinon.SinonStub
 let mockCreateUploadUrl: sinon.SinonStub
@@ -49,7 +48,6 @@ const mockSessionStateConfig = ({
     proxyClient: {
         createConversation: () => sinon.stub(),
         createUploadUrl: () => mockCreateUploadUrl(),
-        generatePlan: () => mockGeneratePlan(),
         startCodeGeneration: () => sinon.stub(),
         getCodeGeneration: () => mockGetCodeGeneration(),
         exportResultArchive: () => mockExportResultArchive(),