fix to remove waitUntilCancellation

Author: ashishrp-aws

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

@@ -1619,15 +1619,7 @@ export class ChatController {
                 return
             }
 
-            await this.messenger.sendAIResponse(
-                response,
-                session,
-                tabID,
-                triggerID,
-                triggerPayload,
-                chatHistory,
-                this.cancelTokenSource.token
-            )
+            await this.messenger.sendAIResponse(response, session, tabID, triggerID, triggerPayload, chatHistory)
         } catch (e: any) {
             this.telemetryHelper.recordMessageResponseError(triggerPayload, tabID, getHttpStatusCode(e) ?? 0)
             // clears session, record telemetry before this call

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

@@ -3,7 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import * as vscode from 'vscode'
 import {
     AppToWebViewMessageDispatcher,
     AuthNeededException,
@@ -69,7 +68,7 @@ import { AsyncEventProgressMessage } from '../../../../amazonq/commons/connector
 import { localize } from '../../../../shared/utilities/vsCodeUtils'
 import { getDiffLinesFromChanges } from '../../../../shared/utilities/diffUtils'
 import { ConversationTracker } from '../../../storages/conversationTracker'
-import { waitUntilWithCancellation } from '../../../../shared/utilities/timeoutUtils'
+import { waitTimeout, Timeout } from '../../../../shared/utilities/timeoutUtils'
 
 export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help'
 
@@ -198,8 +197,7 @@ export class Messenger {
         tabID: string,
         triggerID: string,
         triggerPayload: TriggerPayload,
-        chatHistoryManager: ChatHistoryManager,
-        cancelToken: vscode.CancellationToken
+        chatHistoryManager: ChatHistoryManager
     ) {
         let message = ''
         const messageID = response.$metadata.requestId ?? ''
@@ -237,12 +235,15 @@ export class Messenger {
         })
 
         const eventCounts = new Map<string, number>()
-        await waitUntilWithCancellation(
+        const timeout = new Timeout(600000) // 10 minutes timeout
+
+        await waitTimeout(
             async () => {
                 for await (const chatEvent of response.message!) {
-                    if (cancelToken.isCancellationRequested) {
+                    if (this.isTriggerCancelled(triggerID)) {
                         return
                     }
+
                     for (const key of keys(chatEvent)) {
                         if ((chatEvent[key] as any) !== undefined) {
                             eventCounts.set(key, (eventCounts.get(key) ?? 0) + 1)
@@ -446,11 +447,14 @@ export class Messenger {
                 }
                 return true
             },
+            timeout,
             {
-                timeout: 600000,
-                interval: 500,
-                truthy: true,
-                cancellationToken: conversationTracker.getTokenForTrigger(triggerID) ?? session.tokenSource.token,
+                onCancel: () => {
+                    if (this.isTriggerCancelled(triggerID)) {
+                        return true
+                    }
+                    return false
+                },
             }
         )
             .catch((error: any) => {

packages/core/src/shared/utilities/timeoutUtils.ts

@@ -397,114 +397,3 @@ export function sleep(duration: number = 0): Promise<void> {
     const schedule = globals?.clock?.setTimeout ?? setTimeout
     return new Promise((r) => schedule(r, Math.max(duration, 0)))
 }
-
-/**
- * Similar to waitUntil but with enhanced cancellation support.
- * Waits until the predicate returns true or the operation is cancelled.
- * Continuously checks for cancellation even when waiting for the next chunk.
- *
- * @param fn Function whose result is checked
- * @param options Configuration options including timeout, interval, and cancellation token
- * @returns Result of fn(), or undefined if cancelled
- */
-export async function waitUntilWithCancellation<T>(
-    fn: () => Promise<T>,
-    options: WaitUntilOptions & {
-        cancellationToken: CancellationToken
-    }
-): Promise<T | undefined> {
-    const { cancellationToken, ...waitOptions } = options
-
-    // Set up cancellation listener
-    let cancellationListener: { dispose: () => void } | undefined
-    let checkInterval: NodeJS.Timeout | number | undefined
-    let isCancelled = false
-
-    try {
-        return await new Promise<T | undefined>((resolve, reject) => {
-            // Set up cancellation listener that will resolve with undefined instead of rejecting
-            if (cancellationToken) {
-                cancellationListener = cancellationToken.onCancellationRequested(() => {
-                    if (checkInterval) {
-                        globals.clock.clearInterval(checkInterval)
-                    }
-                    isCancelled = true
-                    resolve(undefined) // Resolve with undefined instead of rejecting
-                })
-            }
-
-            // Use the existing waitUntil function with a wrapper that checks for cancellation
-            waitUntil(
-                async () => {
-                    // Check for cancellation before executing function
-                    if (cancellationToken.isCancellationRequested) {
-                        isCancelled = true
-                        return undefined // Return undefined to signal cancellation
-                    }
-
-                    const result = await fn()
-
-                    // Check for cancellation after executing function
-                    if (cancellationToken.isCancellationRequested) {
-                        isCancelled = true
-                        return undefined // Return undefined to signal cancellation
-                    }
-
-                    return result
-                },
-                {
-                    ...waitOptions,
-                    retryOnFail: (error) => {
-                        // Don't retry if cancelled
-                        if (isCancelled) {
-                            return false
-                        }
-
-                        // Use the original retryOnFail option if provided
-                        if (typeof waitOptions.retryOnFail === 'function') {
-                            return waitOptions.retryOnFail(error)
-                        }
-                        return waitOptions.retryOnFail ?? false
-                    },
-                }
-            )
-                .then((result) => {
-                    // If cancelled during execution, resolve with undefined
-                    if (isCancelled) {
-                        resolve(undefined)
-                    } else {
-                        resolve(result)
-                    }
-                })
-                .catch((error) => {
-                    // If cancelled during execution, resolve with undefined
-                    if (isCancelled) {
-                        resolve(undefined)
-                    } else {
-                        reject(error)
-                    }
-                })
-
-            // Set up an interval to periodically check for cancellation
-            // This ensures we don't miss cancellation events while waiting for the next chunk
-            checkInterval = globals.clock.setInterval(
-                () => {
-                    if (cancellationToken.isCancellationRequested && !isCancelled) {
-                        isCancelled = true
-                        globals.clock.clearInterval(checkInterval)
-                        resolve(undefined) // Resolve with undefined instead of rejecting
-                    }
-                },
-                Math.min(waitOptions.interval ?? waitUntilDefaultInterval, 100)
-            )
-        })
-    } finally {
-        // Clean up resources
-        if (checkInterval) {
-            globals.clock.clearInterval(checkInterval)
-        }
-        if (cancellationListener) {
-            cancellationListener.dispose()
-        }
-    }
-}