fix(amazonqFeatureDev): create a new token to stop and iterate in the same session

Author: tverney

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

@@ -415,7 +415,7 @@ export class FeatureDevController {
             if (session?.state?.tokenSource?.token.isCancellationRequested) {
                 session?.state.tokenSource?.dispose()
                 if (session?.state?.tokenSource) {
-                    session.state.tokenSource = undefined
+                    session.state.tokenSource = new vscode.CancellationTokenSource()
                 }
                 getLogger().debug('Request cancelled, skipping further processing')
             } else {
@@ -712,7 +712,9 @@ export class FeatureDevController {
         })
         this.workOnNewTask(message)
         const session = await this.sessionStorage.getSession(message.tabID)
-        session.state?.tokenSource?.cancel()
+        if (session.state?.tokenSource) {
+            session.state?.tokenSource?.cancel()
+        }
     }
 
     private async tabOpened(message: any) {

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

@@ -26,7 +26,6 @@ import { ReferenceLogViewProvider } from '../../codewhisperer/service/referenceL
 import { AuthUtil } from '../../codewhisperer/util/authUtil'
 import { getLogger } from '../../shared'
 import { logWithConversationId } from '../userFacingText'
-
 export class Session {
     private _state?: SessionState | Omit<SessionState, 'uploadId'>
     private task: string = ''
@@ -136,7 +135,7 @@ export class Session {
 
         if (resp.nextState) {
             // Cancel the request before moving to a new state
-            this.state?.tokenSource?.cancel()
+            if (!this.state?.tokenSource?.token.isCancellationRequested) this.state?.tokenSource?.cancel()
 
             // Move to the next state
             this._state = resp.nextState

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

@@ -39,6 +39,8 @@ import { collectFiles, getWorkspaceFoldersByPrefixes } from '../../shared/utilit
 import { i18n } from '../../shared/i18n-helper'
 import { Messenger } from '../controllers/chat/messenger/messenger'
 
+const EMPTY_CODEGEN_ID = 'EMPTY_CURRENT_CODE_GENERATION_ID'
+
 export class ConversationNotStartedState implements Omit<SessionState, 'uploadId'> {
     public tokenSource: vscode.CancellationTokenSource
     public readonly phase = DevPhase.INIT
@@ -136,7 +138,7 @@ abstract class CodeGenBase {
         this.isCancellationRequested = false
         this.conversationId = config.conversationId
         this.uploadId = config.uploadId
-        this.currentCodeGenerationId = config.currentCodeGenerationId
+        this.currentCodeGenerationId = config.currentCodeGenerationId || EMPTY_CODEGEN_ID
     }
 
     async generateCode({
@@ -160,8 +162,7 @@ abstract class CodeGenBase {
     }> {
         for (
             let pollingIteration = 0;
-            pollingIteration < this.pollCount &&
-            (!this.isCancellationRequested || !this.tokenSource.token.isCancellationRequested);
+            pollingIteration < this.pollCount && !this.isCancellationRequested;
             ++pollingIteration
         ) {
             const codegenResult = await this.config.proxyClient.getCodeGeneration(this.conversationId, codeGenerationId)
@@ -275,8 +276,6 @@ 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
                 })
 
                 action.telemetry.setGenerateCodeIteration(this.currentIteration)
@@ -290,9 +289,6 @@ export class CodeGenState extends CodeGenBase implements SessionState {
                     this.currentCodeGenerationId
                 )
 
-                this.currentCodeGenerationId = codeGenerationId
-                this.config.currentCodeGenerationId = codeGenerationId
-
                 if (!this.isCancellationRequested) {
                     action.messenger.sendAnswer({
                         message: i18n('AWS.amazonq.featureDev.pillText.generatingCode'),
@@ -313,6 +309,11 @@ export class CodeGenState extends CodeGenBase implements SessionState {
                     workspaceFolders: this.config.workspaceFolders,
                 })
 
+                if (codeGeneration && !this.isCancellationRequested) {
+                    this.config.currentCodeGenerationId = codeGenerationId
+                    this.currentCodeGenerationId = codeGenerationId
+                }
+
                 this.filePaths = codeGeneration.newFiles
                 this.deletedFiles = codeGeneration.deletedFiles
                 this.references = codeGeneration.references
@@ -478,7 +479,11 @@ export class PrepareCodeGenState implements SessionState {
                 span
             )
             const uploadId = randomUUID()
-            const { uploadUrl, kmsKeyArn } = await this.config.proxyClient.createUploadUrl(
+            const {
+                uploadUrl,
+                uploadId: returnedUploadId,
+                kmsKeyArn,
+            } = await this.config.proxyClient.createUploadUrl(
                 this.config.conversationId,
                 zipFileChecksum,
                 zipFileBuffer.length,