fix(amazonqFeatureDev): require isCancellation to localize in the ins…

…tance and abort the individual call
aws · Sep 26, 2024 · 0588d8e · 0588d8e
1 parent 132d260
commit 0588d8e
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/packages/core/src/amazonqFeatureDev/session/sessionState.ts b/packages/core/src/amazonqFeatureDev/session/sessionState.ts
@@ -120,6 +120,7 @@ abstract class CodeGenBase {
     public readonly conversationId: string
     public readonly uploadId: string
     public currentCodeGenerationId?: string
+    public isCancellationRequested?: boolean
 
     constructor(
         protected config: SessionStateConfig,
@@ -137,14 +138,12 @@ abstract class CodeGenBase {
         codeGenerationId,
         telemetry: telemetry,
         workspaceFolders,
-        isCancellationRequested,
     }: {
         messenger: Messenger
         fs: VirtualFileSystem
         codeGenerationId: string
         telemetry: TelemetryHelper
         workspaceFolders: CurrentWsFolders
-        isCancellationRequested?: boolean
     }): Promise<{
         newFiles: NewFileInfo[]
         deletedFiles: DeletedFileInfo[]
@@ -154,7 +153,7 @@ abstract class CodeGenBase {
     }> {
         for (
             let pollingIteration = 0;
-            pollingIteration < this.pollCount && !isCancellationRequested;
+            pollingIteration < this.pollCount && !this.isCancellationRequested;
             ++pollingIteration
         ) {
             const codegenResult = await this.config.proxyClient.getCodeGeneration(this.conversationId, codeGenerationId)
@@ -230,7 +229,7 @@ abstract class CodeGenBase {
                 }
             }
         }
-        if (!isCancellationRequested) {
+        if (!this.isCancellationRequested) {
             // still in progress
             const errorMessage = i18n('AWS.amazonq.featureDev.error.codeGen.timeout')
             throw new ToolkitError(errorMessage, { code: 'CodeGenTimeout' })
@@ -260,6 +259,12 @@ export class CodeGenState extends CodeGenBase implements SessionState {
     async interact(action: SessionStateAction): Promise<SessionStateInteraction> {
         return telemetry.amazonq_codeGenerationInvoke.run(async (span) => {
             try {
+                action.tokenSource?.token.onCancellationRequested(() => {
+                    this.isCancellationRequested = true
+                    if (action.tokenSource) {
+                        this.tokenSource = action.tokenSource
+                    }
+                })
                 span.record({
                     amazonqConversationId: this.config.conversationId,
                     credentialStartUrl: AuthUtil.instance.startUrl,
@@ -276,7 +281,7 @@ export class CodeGenState extends CodeGenBase implements SessionState {
                     this.currentCodeGenerationId
                 )
 
-                if (!action.tokenSource?.token.isCancellationRequested) {
+                if (!this.isCancellationRequested) {
                     action.messenger.sendAnswer({
                         message: i18n('AWS.amazonq.featureDev.pillText.generatingCode'),
                         type: 'answer-part',
@@ -294,7 +299,6 @@ export class CodeGenState extends CodeGenBase implements SessionState {
                     codeGenerationId,
                     telemetry: action.telemetry,
                     workspaceFolders: this.config.workspaceFolders,
-                    isCancellationRequested: action.tokenSource?.token.isCancellationRequested,
                 })
 
                 if (codeGeneration && !action.tokenSource?.token.isCancellationRequested) {