perf(amazonq): Adding backoff and retry for /test APIs (#6461)

laileni-aws · web-flow · commit 3a03e95599e5 · 2025-01-29T10:47:23.000-08:00
### Problem - Unit test generation API does not have retry mechanisms. ## Solution - Adding retry mechanisms for UTG APIs for reducing error rate. - No functionality is effected with this PR change, so no new test cases were added. --- - 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.
diff --git a/packages/core/src/codewhisperer/service/securityScanHandler.ts b/packages/core/src/codewhisperer/service/securityScanHandler.ts
@@ -14,7 +14,7 @@ import {
     CodeScanStoppedError,
     onDemandFileScanState,
 } from '../models/model'
-import { sleep } from '../../shared/utilities/timeoutUtils'
+import { sleep, waitUntil } from '../../shared/utilities/timeoutUtils'
 import * as codewhispererClient from '../client/codewhisperer'
 import * as CodeWhispererConstants from '../models/constants'
 import { existsSync, statSync, readFileSync } from 'fs' // eslint-disable-line no-restricted-imports
@@ -363,29 +363,47 @@ export async function uploadArtifactToS3(
         headersObj['x-amz-server-side-encryption-aws-kms-key-id'] = resp.kmsKeyArn
     }
 
-    try {
-        const response = await request.fetch('PUT', resp.uploadUrl, {
-            body: readFileSync(fileName),
-            headers: resp?.requestHeaders ?? headersObj,
-        }).response
-        logger.debug(`StatusCode: ${response.status}, Text: ${response.statusText}`)
-    } catch (error) {
-        let errorMessage = ''
-        const isCodeScan = featureUseCase === FeatureUseCase.CODE_SCAN
-        const featureType = isCodeScan ? 'security scans' : 'unit test generation'
-        const defaultMessage = isCodeScan ? 'Security scan failed.' : 'Test generation failed.'
-        getLogger().error(
-            `Amazon Q is unable to upload workspace artifacts to Amazon S3 for ${featureType}. ` +
-                'For more information, see the Amazon Q documentation or contact your network or organization administrator.'
-        )
-        const errorDesc = getTelemetryReasonDesc(error)
-        if (errorDesc?.includes('"PUT" request failed with code "403"')) {
-            errorMessage = '"PUT" request failed with code "403"'
-        } else if (errorDesc?.includes('"PUT" request failed with code "503"')) {
-            errorMessage = '"PUT" request failed with code "503"'
-        } else {
-            errorMessage = errorDesc ?? defaultMessage
+    let errorMessage = ''
+    const isCodeScan = featureUseCase === FeatureUseCase.CODE_SCAN
+    const result = await waitUntil(
+        async () => {
+            try {
+                const response = await request.fetch('PUT', resp.uploadUrl, {
+                    body: readFileSync(fileName),
+                    headers: resp?.requestHeaders ?? headersObj,
+                }).response
+
+                logger.debug(`StatusCode: ${response.status}, Text: ${response.statusText}`)
+                return response.status === 200 ? response : false
+            } catch (error) {
+                const featureType = isCodeScan ? 'security scans' : 'unit test generation'
+                const defaultMessage = isCodeScan ? 'Security scan failed.' : 'Test generation failed.'
+
+                getLogger().error(
+                    `Amazon Q is unable to upload workspace artifacts to Amazon S3 for ${featureType}. ` +
+                        'For more information, see the Amazon Q documentation or contact your network or organization administrator.'
+                )
+                const errorDesc = getTelemetryReasonDesc(error)
+                const errorCodes = new Map([
+                    ['403', '"PUT" request failed with code "403"'],
+                    ['503', '"PUT" request failed with code "503"'],
+                ])
+
+                errorMessage = errorDesc?.includes('"PUT" request failed with code "')
+                    ? (errorCodes.get(errorDesc.match(/\d+/)?.[0] ?? '') ?? errorDesc)
+                    : (errorDesc ?? defaultMessage)
+
+                return false // Return false to continue retrying
+            }
+        },
+        {
+            interval: 200, // 200ms between attempts
+            timeout: 1000, // 1 second timeout
+            truthy: true,
         }
+    )
+
+    if (!result) {
         throw isCodeScan ? new UploadArtifactToS3Error(errorMessage) : new UploadTestArtifactToS3Error(errorMessage)
     }
 }
diff --git a/packages/core/src/codewhisperer/service/testGenHandler.ts b/packages/core/src/codewhisperer/service/testGenHandler.ts
@@ -23,7 +23,7 @@ import {
     TestGenTimedOutError,
 } from '../../amazonqTest/error'
 import { getMd5, uploadArtifactToS3 } from './securityScanHandler'
-import { fs, randomUUID, sleep, tempDirPath } from '../../shared'
+import { fs, randomUUID, sleep, tempDirPath, waitUntil } from '../../shared'
 import { ShortAnswer, testGenState } from '../models/model'
 import { ChatSessionManager } from '../../amazonqTest/chat/storages/chatSession'
 import { createCodeWhispererChatStreamingClient } from '../../shared/clients/codewhispererChatClient'
@@ -54,17 +54,39 @@ export async function getPresignedUrlAndUploadTestGen(zipMetadata: ZipMetadata)
         uploadIntent: CodeWhispererConstants.testGenUploadIntent,
     }
     logger.verbose(`Prepare for uploading src context...`)
-    const srcResp = await codeWhisperer.codeWhispererClient.createUploadUrl(srcReq).catch((err) => {
-        getLogger().error(`Failed getting presigned url for uploading src context. Request id: ${err.requestId}`)
-        throw new CreateUploadUrlError(err.message)
-    })
-    logger.verbose(`CreateUploadUrlRequest requestId: ${srcResp.$response.requestId}`)
+    let errorMessage = ''
+    const result = await waitUntil(
+        async () => {
+            try {
+                const srcResp = await codeWhisperer.codeWhispererClient.createUploadUrl(srcReq)
+                // Only return the response if it's successful, otherwise return false to continue retrying
+                return srcResp.$response?.httpResponse.statusCode === 200 ? srcResp : false
+            } catch (err: any) {
+                getLogger().error(
+                    `Failed getting presigned url for uploading src context. Request id: ${err.requestId}`
+                )
+                errorMessage = err.message
+                return false // Return false to continue retrying
+            }
+        },
+        {
+            interval: 200, // 200ms between attempts
+            timeout: 1000, // 1 second timeout
+            truthy: true,
+        }
+    )
+
+    if (!result) {
+        throw new CreateUploadUrlError(errorMessage)
+    }
+
+    logger.verbose(`CreateUploadUrlRequest requestId: ${result.$response.requestId}`)
     logger.verbose(`Complete Getting presigned Url for uploading src context.`)
     logger.verbose(`Uploading src context...`)
-    await uploadArtifactToS3(zipMetadata.zipFilePath, srcResp, CodeWhispererConstants.FeatureUseCase.TEST_GENERATION)
+    await uploadArtifactToS3(zipMetadata.zipFilePath, result, CodeWhispererConstants.FeatureUseCase.TEST_GENERATION)
     logger.verbose(`Complete uploading src context.`)
     const artifactMap: ArtifactMap = {
-        SourceCode: srcResp.uploadId,
+        SourceCode: result.uploadId,
     }
     return artifactMap
 }
@@ -102,11 +124,31 @@ export async function createTestJob(
     logger.debug('target line range start: %O', firstTargetLineRangeList?.start)
     logger.debug('target line range end: %O', firstTargetLineRangeList?.end)
 
-    const resp = await codewhispererClient.codeWhispererClient.startTestGeneration(req).catch((err) => {
-        ChatSessionManager.Instance.getSession().startTestGenerationRequestId = err.requestId
-        logger.error(`Failed creating test job. Request id: ${err.requestId}`)
-        throw new CreateTestJobError(err.message)
-    })
+    let errorMessage = ''
+    const resp = await waitUntil(
+        async () => {
+            try {
+                const response = await codewhispererClient.codeWhispererClient.startTestGeneration(req)
+                // Only return the response if it's successful, otherwise return false to continue retrying
+                return response.$response?.httpResponse.statusCode === 200 ? response : false
+            } catch (err: any) {
+                ChatSessionManager.Instance.getSession().startTestGenerationRequestId = err.requestId
+                logger.error(`Failed creating test job. Request id: ${err.requestId}`)
+                errorMessage = err.message
+                return false // Return false to continue retrying
+            }
+        },
+        {
+            interval: 200, // 200ms between attempts
+            timeout: 1000, // 1 second timeout
+            truthy: true,
+        }
+    )
+
+    if (!resp) {
+        throw new CreateTestJobError(errorMessage)
+    }
+
     logger.info('Unit test generation request id: %s', resp.$response.requestId)
     logger.debug('Unit test generation data: %O', resp.$response.data)
     ChatSessionManager.Instance.getSession().startTestGenerationRequestId = resp.$response.requestId
