fix(dev): apply feedback

Author: tverney

packages/core/src/amazonq/webview/ui/connector.ts

@@ -177,8 +177,13 @@ export class Connector {
     }
 
     onStopChatResponse = (tabID: string): void => {
-        if (this.tabsStorage.getTab(tabID)?.type === 'featuredev') {
-            this.featureDevChatConnector.onStopChatResponse(tabID)
+        switch (this.tabsStorage.getTab(tabID)?.type) {
+            case 'featuredev':
+                this.featureDevChatConnector.onStopChatResponse(tabID)
+                break
+            case 'cwc':
+                this.cwChatConnector.onStopChatResponse(tabID)
+                break
         }
     }
 

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

@@ -313,6 +313,15 @@ export class FeatureDevController {
         }
     }
 
+    /**
+     *
+     * This function dispose cancellation token to free resources and provide a new token.
+     * Since user can abort a call in the same session, when the processing ends, we need provide a new one
+     * to start with the new prompt and allow the ability to stop again.
+     *
+     * @param session
+     */
+
     private disposeToken(session: Session | undefined) {
         if (session?.state?.tokenSource?.token.isCancellationRequested) {
             session?.state.tokenSource?.dispose()
@@ -484,7 +493,7 @@ export class FeatureDevController {
     }
     private workOnNewTask(
         message: any,
-        remainingIterations?: number,
+        remainingIterations: number = 0,
         totalIterations?: number,
         isStoppedGeneration: boolean = false
     ) {
@@ -499,7 +508,7 @@ export class FeatureDevController {
             })
         }
 
-        if (((remainingIterations ?? 0) <= 0 && isStoppedGeneration) || !isStoppedGeneration) {
+        if ((remainingIterations <= 0 && isStoppedGeneration) || !isStoppedGeneration) {
             this.messenger.sendAnswer({
                 type: 'system-prompt',
                 tabID: message.tabID,

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

@@ -300,13 +300,6 @@ export class CodeGenState extends CodeGenBase implements SessionState {
                     credentialStartUrl: AuthUtil.instance.startUrl,
                 })
 
-                action.tokenSource?.token.onCancellationRequested(() => {
-                    this.isCancellationRequested = true
-                    if (action.tokenSource) {
-                        this.tokenSource = action.tokenSource
-                    }
-                })
-
                 action.telemetry.setGenerateCodeIteration(this.currentIteration)
                 action.telemetry.setGenerateCodeLastInvocationTime()
                 const codeGenerationId = randomUUID()
@@ -496,7 +489,7 @@ export class PrepareCodeGenState implements SessionState {
         public codeGenerationRemainingIterationCount?: number,
         public codeGenerationTotalIterationCount?: number,
         public uploadHistory: UploadHistory = {},
-        public superTokenSource?: vscode.CancellationTokenSource,
+        public superTokenSource: vscode.CancellationTokenSource = new vscode.CancellationTokenSource(),
         public currentCodeGenerationId?: string,
         public codeGenerationId?: string
     ) {