refactor(timeout): make retryOnFail always a function

packages/amazonq/test/e2e/inline/inline.test.ts

@@ -24,7 +24,7 @@ describe('Amazon Q Inline', async function () {
     const waitOptions = {
         interval: 500,
         timeout: 10000,
-        retryOnFail: false,
+        retryOnFail: () => true,
     }
 
     before(async function () {

packages/core/src/notifications/controller.ts

@@ -298,7 +298,7 @@ export class RemoteFetcher implements NotificationFetcher {
             {
                 interval: RemoteFetcher.retryIntervalMs,
                 timeout: RemoteFetcher.retryTimeout,
-                retryOnFail: true,
+                retryOnFail: () => true,
                 // No exponential backoff - necessary?
             }
         )

packages/core/src/shared/crashMonitoring.ts

@@ -492,7 +492,7 @@ export class FileSystemState {
                 timeout: 15_000,
                 interval: 100,
                 backoff: 2,
-                retryOnFail: true,
+                retryOnFail: () => true,
             })
 
             return funcWithRetries
@@ -550,7 +550,7 @@ export class FileSystemState {
             timeout: 15_000,
             interval: 100,
             backoff: 2,
-            retryOnFail: true,
+            retryOnFail: () => true,
         })
         await funcWithRetries
     }
@@ -618,7 +618,12 @@ export class FileSystemState {
                 }
                 const funcWithIgnoreBadFile = () => ignoreBadFileError(loadExtFromDisk)
                 const funcWithRetries = () =>
-                    waitUntil(funcWithIgnoreBadFile, { timeout: 15_000, interval: 100, backoff: 2, retryOnFail: true })
+                    waitUntil(funcWithIgnoreBadFile, {
+                        timeout: 15_000,
+                        interval: 100,
+                        backoff: 2,
+                        retryOnFail: () => true,
+                    })
                 const funcWithCtx = () => withFailCtx('parseRunningExtFile', funcWithRetries)
                 const ext: ExtInstanceHeartbeat | undefined = await funcWithCtx()
 

packages/core/src/shared/resourcefetcher/httpResourceFetcher.ts

@@ -120,9 +120,9 @@ export class HttpResourceFetcher implements ResourceFetcher<Response> {
                 timeout: 3000,
                 interval: 100,
                 backoff: 2,
-                retryOnFail: (error: Error) => {
-                    // Retry unless the user intentionally canceled the operation.
-                    return !isUserCancelledError(error)
+                retryOnFail: (error: Error | undefined) => {
+                    // Only stops retry when the user intentionally canceled the operation.
+                    return error ? !isUserCancelledError(error) : true
                 },
             }
         )

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

@@ -222,13 +222,16 @@ interface WaitUntilOptions {
     /** A backoff multiplier for how long the next interval will be (default: None, i.e 1) */
     readonly backoff?: number
     /**
-     * Only retries when an error is thrown, otherwise returning the immediate result.
-     * Can also be a callback for conditional retry based on errors
-     * - 'truthy' arg is ignored
-     * - If the timeout is reached it throws the last error
-     * - default: false
+     * Call back, true when the function should be retried on failure.
+     * Example usage:
+     * - () => boolean if the retry logic is "constant"
+     * - (error) => { error ? retryOnError(error) : boolean }
+     *   where retryOnError determines what errors to retry, and the boolean is for base case when there's no error (undefined)
+     * 'truthy' arg is ignored
+     * If the timeout is reached it throws the last error
+     * Default to false
      */
-    readonly retryOnFail?: boolean | ((error: Error) => boolean)
+    readonly retryOnFail?: (error: Error | undefined) => boolean
 }
 
 export const waitUntilDefaultTimeout = 2000
@@ -243,14 +246,9 @@ export const waitUntilDefaultInterval = 500
  *
  * @returns Result of `fn()`, or possibly `undefined` depending on the arguments.
  */
-export async function waitUntil<T>(fn: () => Promise<T>, options: WaitUntilOptions & { retryOnFail: true }): Promise<T>
 export async function waitUntil<T>(
     fn: () => Promise<T>,
-    options: WaitUntilOptions & { retryOnFail: false }
-): Promise<T | undefined>
-export async function waitUntil<T>(
-    fn: () => Promise<T>,
-    options: WaitUntilOptions & { retryOnFail: (error: Error) => boolean }
+    options: WaitUntilOptions & { retryOnFail: (error: Error | undefined) => boolean }
 ): Promise<T>
 
 export async function waitUntil<T>(
@@ -264,7 +262,7 @@ export async function waitUntil<T>(fn: () => Promise<T>, options: WaitUntilOptio
         interval: waitUntilDefaultInterval,
         truthy: true,
         backoff: 1,
-        retryOnFail: false,
+        retryOnFail: () => false,
         ...options,
     }
 
@@ -273,17 +271,6 @@ export async function waitUntil<T>(fn: () => Promise<T>, options: WaitUntilOptio
     let elapsed: number = 0
     let remaining = opt.timeout
 
-    // Internal helper to determine if we should retry
-    function shouldRetry(error: Error | undefined): boolean {
-        if (error === undefined) {
-            return typeof opt.retryOnFail === 'boolean' ? opt.retryOnFail : true
-        }
-        if (typeof opt.retryOnFail === 'function') {
-            return opt.retryOnFail(error)
-        }
-        return opt.retryOnFail
-    }
-
     for (let i = 0; true; i++) {
         const start: number = globals.clock.Date.now()
         let result: T
@@ -296,7 +283,7 @@ export async function waitUntil<T>(fn: () => Promise<T>, options: WaitUntilOptio
                 result = await fn()
             }
 
-            if (shouldRetry(lastError) || (opt.truthy && result) || (!opt.truthy && result !== undefined)) {
+            if (opt.retryOnFail(lastError) || (opt.truthy && result) || (!opt.truthy && result !== undefined)) {
                 return result
             }
         } catch (e) {
@@ -305,7 +292,7 @@ export async function waitUntil<T>(fn: () => Promise<T>, options: WaitUntilOptio
                 throw e
             }
 
-            if (!shouldRetry(e)) {
+            if (!opt.retryOnFail(e)) {
                 throw e
             }
 
@@ -317,7 +304,7 @@ export async function waitUntil<T>(fn: () => Promise<T>, options: WaitUntilOptio
 
         // If the sleep will exceed the timeout, abort early
         if (elapsed + interval >= remaining) {
-            if (!shouldRetry(lastError)) {
+            if (!opt.retryOnFail(lastError)) {
                 return undefined
             }
             throw lastError

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

@@ -400,8 +400,8 @@ export const timeoutUtilsDescribe = describe('timeoutUtils', async function () {
             fn.onCall(2).resolves('success')
 
             const res = waitUntil(fn, {
-                retryOnFail: (error: Error) => {
-                    return error.message === 'Retry error'
+                retryOnFail: (error) => {
+                    return error ? error.message === 'Retry error' : false
                 },
             })
 
@@ -420,8 +420,8 @@ export const timeoutUtilsDescribe = describe('timeoutUtils', async function () {
 
             const res = assert.rejects(
                 waitUntil(fn, {
-                    retryOnFail: (error: Error) => {
-                        return error.message === 'Retry error'
+                    retryOnFail: (error) => {
+                        return error ? error.message === 'Retry error' : true
                     },
                 }),
                 (e) => e instanceof Error && e.message === 'Last error'
@@ -439,7 +439,7 @@ export const timeoutUtilsDescribe = describe('timeoutUtils', async function () {
             fn.onCall(1).throws()
             fn.onCall(2).resolves('success')
 
-            const res = waitUntil(fn, { retryOnFail: true })
+            const res = waitUntil(fn, { retryOnFail: () => true })
 
             await clock.tickAsync(timeoutUtils.waitUntilDefaultInterval)
             assert.strictEqual(fn.callCount, 2)
@@ -450,7 +450,7 @@ export const timeoutUtilsDescribe = describe('timeoutUtils', async function () {
 
         it('retryOnFail ignores truthiness', async function () {
             fn.resolves(false)
-            const res = waitUntil(fn, { retryOnFail: true, truthy: true })
+            const res = waitUntil(fn, { retryOnFail: () => true, truthy: true })
             assert.strictEqual(await res, false)
         })
 
@@ -466,7 +466,7 @@ export const timeoutUtilsDescribe = describe('timeoutUtils', async function () {
             // We must wrap w/ assert.rejects() here instead of at the end, otherwise Mocha raise a
             // `rejected promise not handled within 1 second: Error: last`
             const res = assert.rejects(
-                waitUntil(fn, { retryOnFail: true }),
+                waitUntil(fn, { retryOnFail: () => true }),
                 (e) => e instanceof Error && e.message === 'last'
             )
 
@@ -488,7 +488,7 @@ export const timeoutUtilsDescribe = describe('timeoutUtils', async function () {
 
             // Note 701 instead of 700 for timeout. The 1 millisecond allows the final call to execute
             // since the timeout condition is >= instead of >
-            const res = waitUntil(fn, { timeout: 701, interval: 100, backoff: 2, retryOnFail: true })
+            const res = waitUntil(fn, { timeout: 701, interval: 100, backoff: 2, retryOnFail: () => true })
 
             // Check the call count after each iteration, ensuring the function is called
             // after the correct delay between retries.