feat(amazonq): user cancellable lsp download (#6573)

tomcat323 · tomzu · web-flow · commit ff1bbbd912e5 · 2025-03-04T09:42:35.000-05:00
## Problem the cancellation button on notification doesn't actually stop the LSP download ## Solution Route the notification timeout into `httpResourceFetcher` call, so that the user cancellation event stops the download as well. Demo: non-cancelled and cancelled downloads: https://github.com/user-attachments/assets/4f3d4ed0-635a-4f88-aef7-ee62e87bfc83 --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license. --------- Co-authored-by: tomzu <tomzu@amazon.com>
diff --git a/packages/core/src/shared/errors.ts b/packages/core/src/shared/errors.ts
@@ -16,6 +16,7 @@ import { CodeWhispererStreamingServiceException } from '@amzn/codewhisperer-stre
 import { driveLetterRegex } from './utilities/pathUtils'
 import { getLogger } from './logger/logger'
 import { crashMonitoringDirName } from './constants'
+import { RequestCancelledError } from './request'
 
 let _username = 'unknown-user'
 let _isAutomation = false
@@ -618,7 +619,11 @@ function hasTime(error: Error): error is typeof error & { time: Date } {
 }
 
 export function isUserCancelledError(error: unknown): boolean {
-    return CancellationError.isUserCancelled(error) || (error instanceof ToolkitError && error.cancelled)
+    return (
+        CancellationError.isUserCancelled(error) ||
+        (error instanceof ToolkitError && error.cancelled) ||
+        error instanceof RequestCancelledError
+    )
 }
 
 /**
diff --git a/packages/core/src/shared/lsp/lspResolver.ts b/packages/core/src/shared/lsp/lspResolver.ts
@@ -16,6 +16,9 @@ import { HttpResourceFetcher } from '../resourcefetcher/httpResourceFetcher'
 import { showMessageWithCancel } from '../../shared/utilities/messages'
 import { Timeout } from '../utilities/timeoutUtils'
 
+// max timeout for downloading remote LSP assets progress, the lowest possible is 3000, bounded by httpResourceFetcher's waitUntil
+const remoteDownloadTimeout = 5000
+
 export class LanguageServerResolver {
     constructor(
         private readonly manifest: Manifest,
@@ -32,8 +35,6 @@ export class LanguageServerResolver {
      * @throws ToolkitError if no compatible version can be found
      */
     async resolve() {
-        const timeout = new Timeout(5000)
-        await showMessageWithCancel(`Downloading '${this.lsName}' language server`, timeout)
         function getServerVersion(result: LspResult) {
             return {
                 languageServerVersion: result.version,
@@ -62,7 +63,6 @@ export class LanguageServerResolver {
             return await tryStageResolvers('getServer', serverResolvers, getServerVersion)
         } finally {
             logger.info(`Finished setting up LSP server`)
-            timeout.cancel()
         }
     }
 
@@ -86,19 +86,34 @@ export class LanguageServerResolver {
         }
     }
 
+    /**
+     * Show a toast notification with progress bar for lsp remote downlaod
+     * Returns a timeout to be passed down into httpFetcher to handle user cancellation
+     */
+    private async showDownloadProgress() {
+        const timeout = new Timeout(remoteDownloadTimeout)
+        await showMessageWithCancel(`Downloading '${this.lsName}' language server`, timeout)
+        return timeout
+    }
+
     private async fetchRemoteServer(
         cacheDirectory: string,
         latestVersion: LspVersion,
         targetContents: TargetContent[]
     ): Promise<LspResult> {
-        if (await this.downloadRemoteTargetContent(targetContents, latestVersion.serverVersion)) {
-            return {
-                location: 'remote',
-                version: latestVersion.serverVersion,
-                assetDirectory: cacheDirectory,
+        const timeout = await this.showDownloadProgress()
+        try {
+            if (await this.downloadRemoteTargetContent(targetContents, latestVersion.serverVersion, timeout)) {
+                return {
+                    location: 'remote',
+                    version: latestVersion.serverVersion,
+                    assetDirectory: cacheDirectory,
+                }
+            } else {
+                throw new ToolkitError('Failed to download server from remote', { code: 'RemoteDownloadFailed' })
             }
-        } else {
-            throw new ToolkitError('Failed to download server from remote', { code: 'RemoteDownloadFailed' })
+        } finally {
+            timeout.dispose()
         }
     }
 
@@ -192,7 +207,7 @@ export class LanguageServerResolver {
      *  true, if all of the contents were successfully downloaded and unzipped
      *  false, if any of the contents failed to download or unzip
      */
-    private async downloadRemoteTargetContent(contents: TargetContent[], version: string) {
+    private async downloadRemoteTargetContent(contents: TargetContent[], version: string, timeout: Timeout) {
         const downloadDirectory = this.getDownloadDirectory(version)
 
         if (!(await fs.existsDir(downloadDirectory))) {
@@ -201,7 +216,11 @@ export class LanguageServerResolver {
 
         const fetchTasks = contents.map(async (content) => {
             return {
-                res: await new HttpResourceFetcher(content.url, { showUrl: true }).get(),
+                res: await new HttpResourceFetcher(content.url, {
+                    showUrl: true,
+                    timeout: timeout,
+                    throwOnError: true,
+                }).get(),
                 hash: content.hashes[0],
                 filename: content.filename,
             }
diff --git a/packages/core/src/shared/resourcefetcher/httpResourceFetcher.ts b/packages/core/src/shared/resourcefetcher/httpResourceFetcher.ts
@@ -8,6 +8,7 @@ import { getLogger, Logger } from '../logger/logger'
 import { ResourceFetcher } from './resourcefetcher'
 import { Timeout, CancelEvent, waitUntil } from '../utilities/timeoutUtils'
 import request, { RequestError } from '../request'
+import { isUserCancelledError } from '../errors'
 
 type RequestHeaders = { eTag?: string; gZip?: boolean }
 
@@ -22,15 +23,19 @@ export class HttpResourceFetcher implements ResourceFetcher<Response> {
      * @param {string} params.friendlyName If URL is not shown, replaces the URL with this text.
      * @param {Timeout} params.timeout Timeout token to abort/cancel the request. Similar to `AbortSignal`.
      * @param {number} params.retries The number of retries a get request should make if one fails
+     * @param {boolean} params.throwOnError True if we want to throw if there's request error, defaul to false
      */
     public constructor(
         private readonly url: string,
         private readonly params: {
             showUrl: boolean
             friendlyName?: string
             timeout?: Timeout
+            throwOnError?: boolean
         }
-    ) {}
+    ) {
+        this.params.throwOnError = this.params.throwOnError ?? false
+    }
 
     /**
      * Returns the response of the resource, or undefined if the response failed could not be retrieved.
@@ -82,6 +87,9 @@ export class HttpResourceFetcher implements ResourceFetcher<Response> {
                 `Error downloading ${this.logText()}: %s`,
                 error.message ?? error.code ?? error.response.statusText ?? error.response.status
             )
+            if (this.params.throwOnError) {
+                throw error
+            }
             return undefined
         }
     }
@@ -112,7 +120,10 @@ export class HttpResourceFetcher implements ResourceFetcher<Response> {
                 timeout: 3000,
                 interval: 100,
                 backoff: 2,
-                retryOnFail: true,
+                retryOnFail: (error: Error) => {
+                    // Retry unless the user intentionally canceled the operation.
+                    return !isUserCancelledError(error)
+                },
             }
         )
     }
diff --git a/packages/core/src/shared/utilities/timeoutUtils.ts b/packages/core/src/shared/utilities/timeoutUtils.ts
@@ -223,11 +223,12 @@ interface WaitUntilOptions {
     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
      */
-    readonly retryOnFail?: boolean
+    readonly retryOnFail?: boolean | ((error: Error) => boolean)
 }
 
 export const waitUntilDefaultTimeout = 2000
@@ -247,6 +248,11 @@ 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 }
+): Promise<T>
+
 export async function waitUntil<T>(
     fn: () => Promise<T>,
     options: Omit<WaitUntilOptions, 'retryOnFail'>
@@ -267,6 +273,17 @@ 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
@@ -279,16 +296,16 @@ export async function waitUntil<T>(fn: () => Promise<T>, options: WaitUntilOptio
                 result = await fn()
             }
 
-            if (opt.retryOnFail || (opt.truthy && result) || (!opt.truthy && result !== undefined)) {
+            if (shouldRetry(lastError) || (opt.truthy && result) || (!opt.truthy && result !== undefined)) {
                 return result
             }
         } catch (e) {
-            if (!opt.retryOnFail) {
+            // Unlikely to hit this, but exists for typing
+            if (!(e instanceof Error)) {
                 throw e
             }
 
-            // Unlikely to hit this, but exists for typing
-            if (!(e instanceof Error)) {
+            if (!shouldRetry(e)) {
                 throw e
             }
 
@@ -300,10 +317,9 @@ export async function waitUntil<T>(fn: () => Promise<T>, options: WaitUntilOptio
 
         // If the sleep will exceed the timeout, abort early
         if (elapsed + interval >= remaining) {
-            if (!opt.retryOnFail) {
+            if (!shouldRetry(lastError)) {
                 return undefined
             }
-
             throw lastError
         }
 
diff --git a/packages/core/src/test/shared/utilities/timeoutUtils.test.ts b/packages/core/src/test/shared/utilities/timeoutUtils.test.ts
@@ -394,6 +394,46 @@ export const timeoutUtilsDescribe = describe('timeoutUtils', async function () {
             fn = sandbox.stub()
         })
 
+        it('should retry when retryOnFail callback returns true', async function () {
+            fn.onCall(0).throws(new Error('Retry error'))
+            fn.onCall(1).throws(new Error('Retry error'))
+            fn.onCall(2).resolves('success')
+
+            const res = waitUntil(fn, {
+                retryOnFail: (error: Error) => {
+                    return error.message === 'Retry error'
+                },
+            })
+
+            await clock.tickAsync(timeoutUtils.waitUntilDefaultInterval)
+            assert.strictEqual(fn.callCount, 2)
+            await clock.tickAsync(timeoutUtils.waitUntilDefaultInterval)
+            assert.strictEqual(fn.callCount, 3)
+            assert.strictEqual(await res, 'success')
+        })
+
+        it('should not retry when retryOnFail callback returns false', async function () {
+            fn.onCall(0).throws(new Error('Retry error'))
+            fn.onCall(1).throws(new Error('Retry error'))
+            fn.onCall(2).throws(new Error('Last error'))
+            fn.onCall(3).resolves('this is not hit')
+
+            const res = assert.rejects(
+                waitUntil(fn, {
+                    retryOnFail: (error: Error) => {
+                        return error.message === 'Retry error'
+                    },
+                }),
+                (e) => e instanceof Error && e.message === 'Last error'
+            )
+
+            await clock.tickAsync(timeoutUtils.waitUntilDefaultInterval)
+            assert.strictEqual(fn.callCount, 2)
+            await clock.tickAsync(timeoutUtils.waitUntilDefaultInterval)
+            assert.strictEqual(fn.callCount, 3)
+            await res
+        })
+
         it('retries the function until it succeeds', async function () {
             fn.onCall(0).throws()
             fn.onCall(1).throws()
