UTG Error handling rebasing

aws · Jan 16, 2025 · 6c25777 · 6c25777
1 parent 62ef0e5
commit 6c25777
Show file tree

Hide file tree

Showing 9 changed files with 201 additions and 83 deletions.
diff --git a/packages/core/src/amazonqTest/chat/controller/controller.ts b/packages/core/src/amazonqTest/chat/controller/controller.ts
@@ -20,6 +20,7 @@ import {
     TelemetryHelper,
     TestGenerationBuildStep,
     testGenState,
+    tooManyRequestErrorMessage,
     unitTestGenerationCancelMessage,
     UserWrittenCodeTracker,
 } from '../../../codewhisperer'
@@ -242,72 +243,76 @@ export class TestController {
         // eslint-disable-next-line unicorn/no-null
         this.messenger.sendUpdatePromptProgress(data.tabID, null)
         const session = this.sessionStorage.getSession()
-        const isCancel = data.error.message === unitTestGenerationCancelMessage
-
+        const isCancel = data.error.uiMessage === unitTestGenerationCancelMessage
+        let telemetryErrorMessage = getTelemetryReasonDesc(data.error)
+        if (session.stopIteration) {
+            telemetryErrorMessage = getTelemetryReasonDesc(data.error.uiMessage.replaceAll('```', ''))
+        }
         TelemetryHelper.instance.sendTestGenerationToolkitEvent(
             session,
             true,
             true,
             isCancel ? 'Cancelled' : 'Failed',
             session.startTestGenerationRequestId,
             performance.now() - session.testGenerationStartTime,
-            getTelemetryReasonDesc(data.error),
+            telemetryErrorMessage,
             session.isCodeBlockSelected,
             session.artifactsUploadDuration,
             session.srcPayloadSize,
             session.srcZipFileSize
         )
-
         if (session.stopIteration) {
             // Error from Science
-            this.messenger.sendMessage(data.error.message.replaceAll('```', ''), data.tabID, 'answer')
+            this.messenger.sendMessage(data.error.uiMessage.replaceAll('```', ''), data.tabID, 'answer')
         } else {
             isCancel
-                ? this.messenger.sendMessage(data.error.message, data.tabID, 'answer')
+                ? this.messenger.sendMessage(data.error.uiMessage, data.tabID, 'answer')
                 : this.sendErrorMessage(data)
         }
         await this.sessionCleanUp()
         return
     }
     // Client side error messages
-    private sendErrorMessage(data: { tabID: string; error: { code: string; message: string } }) {
+    private sendErrorMessage(data: {
+        tabID: string
+        error: { uiMessage: string; message: string; code: string; statusCode: string }
+    }) {
         const { error, tabID } = data
 
+        // If user reached monthly limit for builderId
+        if (error.code === 'CreateTestJobError') {
+            if (error.message.includes(CodeWhispererConstants.utgLimitReached)) {
+                getLogger().error('Monthly quota reached for QSDA actions.')
+                return this.messenger.sendMessage(
+                    i18n('AWS.amazonq.featureDev.error.monthlyLimitReached'),
+                    tabID,
+                    'answer'
+                )
+            }
+            if (error.message.includes('Too many requests')) {
+                getLogger().error(error.message)
+                return this.messenger.sendErrorMessage(tooManyRequestErrorMessage, tabID)
+            }
+        }
         if (isAwsError(error)) {
             if (error.code === 'ThrottlingException') {
-                // TODO: use the explicitly modeled exception reason for quota vs throttle
-                if (error.message.includes(CodeWhispererConstants.utgLimitReached)) {
-                    getLogger().error('Monthly quota reached for QSDA actions.')
-                    return this.messenger.sendMessage(
-                        i18n('AWS.amazonq.featureDev.error.monthlyLimitReached'),
-                        tabID,
-                        'answer'
-                    )
-                } else {
-                    getLogger().error('Too many requests.')
-                    // TODO: move to constants file
-                    this.messenger.sendErrorMessage('Too many requests. Please wait before retrying.', tabID)
-                }
-            } else {
-                // other service errors:
-                // AccessDeniedException - should not happen because access is validated before this point in the client
-                // ValidationException - shouldn't happen because client should not send malformed requests
-                // ConflictException - should not happen because the client will maintain proper state
-                // InternalServerException - shouldn't happen but needs to be caught
-                getLogger().error('Other error message: %s', error.message)
-                this.messenger.sendErrorMessage(
-                    'Encountered an unexpected error when generating tests. Please try again',
-                    tabID
-                )
+                // TODO: use the explicitly modeled exception reason for quota vs throttle{
+                getLogger().error(error.message)
+                this.messenger.sendErrorMessage(tooManyRequestErrorMessage, tabID)
+                return
             }
-        } else {
-            // other unexpected errors (TODO enumerate all other failure cases)
+            // other service errors:
+            // AccessDeniedException - should not happen because access is validated before this point in the client
+            // ValidationException - shouldn't happen because client should not send malformed requests
+            // ConflictException - should not happen because the client will maintain proper state
+            // InternalServerException - shouldn't happen but needs to be caught
             getLogger().error('Other error message: %s', error.message)
-            this.messenger.sendErrorMessage(
-                'Encountered an unexpected error when generating tests. Please try again',
-                tabID
-            )
+            this.messenger.sendErrorMessage('', tabID)
+            return
         }
+        // other unexpected errors (TODO enumerate all other failure cases)
+        getLogger().error('Other error message: %s', error.uiMessage)
+        this.messenger.sendErrorMessage('', tabID)
     }
 
     // This function handles actions if user clicked on any Button one of these cases will be executed
@@ -730,6 +735,9 @@ export class TestController {
         // this.messenger.sendMessage('Accepted', message.tabID, 'prompt')
         telemetry.ui_click.emit({ elementId: 'unitTestGeneration_acceptDiff' })
 
+        getLogger().info(
+            `Generated unit tests are accepted for ${session.fileLanguage ?? 'plaintext'} language with jobId: ${session.listOfTestGenerationJobId[0]}, jobGroupName: ${session.testGenerationJobGroupName}, result: Succeeded`
+        )
         TelemetryHelper.instance.sendTestGenerationToolkitEvent(
             session,
             true,
@@ -751,7 +759,6 @@ export class TestController {
         )
 
         await this.endSession(message, FollowUpTypes.SkipBuildAndFinish)
-        await this.sessionCleanUp()
         return
 
         if (session.listOfTestGenerationJobId.length === 1) {
@@ -876,16 +883,12 @@ export class TestController {
                 session.numberOfTestsGenerated,
                 session.linesOfCodeGenerated
             )
-
             telemetry.ui_click.emit({ elementId: 'unitTestGeneration_rejectDiff' })
         }
 
         await this.sessionCleanUp()
-        // TODO: revert 'Accepted' to 'Skip build and finish' once supported
-        const message = step === FollowUpTypes.RejectCode ? 'Rejected' : 'Accepted'
 
-        this.messenger.sendMessage(message, data.tabID, 'prompt')
-        this.messenger.sendMessage(`Unit test generation workflow is completed.`, data.tabID, 'answer')
+        // this.messenger.sendMessage(`Unit test generation workflow is completed.`, data.tabID, 'answer')
         this.messenger.sendChatInputEnabled(data.tabID, true)
         return
     }
@@ -1320,8 +1323,18 @@ export class TestController {
             'Deleting output.log and temp result directory. testGenerationLogsDir: %s',
             testGenerationLogsDir
         )
-        await fs.delete(path.join(testGenerationLogsDir, 'output.log'))
-        await fs.delete(this.tempResultDirPath, { recursive: true })
+        const outputLogPath = path.join(testGenerationLogsDir, 'output.log')
+        if (await fs.existsFile(outputLogPath)) {
+            await fs.delete(outputLogPath)
+        }
+        if (
+            await fs
+                .stat(this.tempResultDirPath)
+                .then(() => true)
+                .catch(() => false)
+        ) {
+            await fs.delete(this.tempResultDirPath, { recursive: true })
+        }
     }
 
     // TODO: return build command when product approves

diff --git a/packages/core/src/amazonqTest/chat/controller/messenger/messenger.ts b/packages/core/src/amazonqTest/chat/controller/messenger/messenger.ts
@@ -282,9 +282,21 @@ export class Messenger {
                         'Cancelled',
                         messageId,
                         performance.now() - session.testGenerationStartTime,
-                        getTelemetryReasonDesc(CodeWhispererConstants.unitTestGenerationCancelMessage)
+                        getTelemetryReasonDesc(
+                            `TestGenCancelled: ${CodeWhispererConstants.unitTestGenerationCancelMessage}`
+                        ),
+                        undefined,
+                        undefined,
+                        undefined,
+                        undefined,
+                        undefined,
+                        undefined,
+                        undefined,
+                        undefined,
+                        undefined,
+                        undefined,
+                        'TestGenCancelled'
                     )
-
                     this.dispatcher.sendUpdatePromptProgress(
                         new UpdatePromptProgressMessage(tabID, cancellingProgressField)
                     )
@@ -296,9 +308,9 @@ export class Messenger {
                         fileInWorkspace,
                         'Succeeded',
                         messageId,
-                        performance.now() - session.testGenerationStartTime
+                        performance.now() - session.testGenerationStartTime,
+                        undefined
                     )
-
                     this.dispatcher.sendUpdatePromptProgress(
                         new UpdatePromptProgressMessage(tabID, testGenCompletedField)
                     )

diff --git a/packages/core/src/amazonqTest/error.ts b/packages/core/src/amazonqTest/error.ts
@@ -0,0 +1,67 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import { ToolkitError } from '../shared/errors'
+
+export const technicalErrorCustomerFacingMessage =
+    'I am experiencing technical difficulties at the moment. Please try again in a few minutes.'
+const defaultTestGenErrorMessage = 'Amazon Q encountered an error while generating tests. Try again later.'
+export class TestGenError extends ToolkitError {
+    constructor(
+        error: string,
+        code: string,
+        public uiMessage: string
+    ) {
+        super(error, { code })
+    }
+}
+export class ProjectZipError extends TestGenError {
+    constructor(error: string) {
+        super(error, 'ProjectZipError', defaultTestGenErrorMessage)
+    }
+}
+export class InvalidSourceZipError extends TestGenError {
+    constructor() {
+        super('Failed to create valid source zip', 'InvalidSourceZipError', defaultTestGenErrorMessage)
+    }
+}
+export class CreateUploadUrlError extends TestGenError {
+    constructor(errorMessage: string) {
+        super(errorMessage, 'CreateUploadUrlError', technicalErrorCustomerFacingMessage)
+    }
+}
+export class UploadTestArtifactToS3Error extends TestGenError {
+    constructor(error: string) {
+        super(error, 'UploadTestArtifactToS3Error', technicalErrorCustomerFacingMessage)
+    }
+}
+export class CreateTestJobError extends TestGenError {
+    constructor(error: string) {
+        super(error, 'CreateTestJobError', technicalErrorCustomerFacingMessage)
+    }
+}
+export class TestGenTimedOutError extends TestGenError {
+    constructor() {
+        super(
+            'Test generation failed. Amazon Q timed out.',
+            'TestGenTimedOutError',
+            technicalErrorCustomerFacingMessage
+        )
+    }
+}
+export class TestGenStoppedError extends TestGenError {
+    constructor() {
+        super('Unit test generation cancelled.', 'TestGenCancelled', 'Unit test generation cancelled.')
+    }
+}
+export class TestGenFailedError extends TestGenError {
+    constructor(error?: string) {
+        super(error ?? 'Test generation failed', 'TestGenFailedError', error ?? technicalErrorCustomerFacingMessage)
+    }
+}
+export class ExportResultsArchiveError extends TestGenError {
+    constructor(error?: string) {
+        super(error ?? 'Test generation failed', 'ExportResultsArchiveError', technicalErrorCustomerFacingMessage)
+    }
+}
diff --git a/packages/core/src/codewhisperer/commands/startTestGeneration.ts b/packages/core/src/codewhisperer/commands/startTestGeneration.ts
@@ -21,7 +21,7 @@ import { ChildProcess, spawn } from 'child_process' // eslint-disable-line no-re
 import { BuildStatus } from '../../amazonqTest/chat/session/session'
 import { fs } from '../../shared/fs/fs'
 import { TestGenerationJobStatus } from '../models/constants'
-import { TestGenFailedError } from '../models/errors'
+import { TestGenFailedError } from '../../amazonqTest/error'
 import { Range } from '../client/codewhispereruserclient'
 
 // eslint-disable-next-line unicorn/no-null
@@ -75,8 +75,9 @@ export async function startTestGenerationProcess(
         try {
             artifactMap = await getPresignedUrlAndUploadTestGen(zipMetadata)
         } finally {
-            if (await fs.existsFile(path.join(testGenerationLogsDir, 'output.log'))) {
-                await fs.delete(path.join(testGenerationLogsDir, 'output.log'))
+            const outputLogPath = path.join(testGenerationLogsDir, 'output.log')
+            if (await fs.existsFile(outputLogPath)) {
+                await fs.delete(outputLogPath)
             }
             await zipUtil.removeTmpFiles(zipMetadata)
             session.artifactsUploadDuration = performance.now() - uploadStartTime

diff --git a/packages/core/src/codewhisperer/models/constants.ts b/packages/core/src/codewhisperer/models/constants.ts
@@ -726,6 +726,8 @@ export const noOpenProjectsFoundChatTestGenMessage = `Sorry, I couldn\'t find a
 
 export const unitTestGenerationCancelMessage = 'Unit test generation cancelled.'
 
+export const tooManyRequestErrorMessage = 'Too many requests. Please wait before retrying.'
+
 export const noJavaProjectsFoundChatMessage = `I couldn\'t find a project that I can upgrade. Currently, I support Java 8, Java 11, and Java 17 projects built on Maven. Make sure your project is open in the IDE. For more information, see the [Amazon Q documentation](${codeTransformPrereqDoc}).`
 
 export const linkToDocsHome = 'https://docs.aws.amazon.com/amazonq/latest/aws-builder-use-ug/code-transformation.html'
@@ -861,7 +863,7 @@ export enum TestGenerationJobStatus {
     COMPLETED = 'COMPLETED',
 }
 
-export enum ZipUseCase {
+export enum FeatureUseCase {
     TEST_GENERATION = 'TEST_GENERATION',
     CODE_SCAN = 'CODE_SCAN',
 }

diff --git a/packages/core/src/codewhisperer/service/codeFixHandler.ts b/packages/core/src/codewhisperer/service/codeFixHandler.ts
@@ -34,7 +34,7 @@ export async function getPresignedUrlAndUpload(
     getLogger().verbose(`CreateUploadUrlRequest requestId: ${srcResp.$response.requestId}`)
     getLogger().verbose(`Complete Getting presigned Url for uploading src context.`)
     getLogger().verbose(`Uploading src context...`)
-    await uploadArtifactToS3(zipFilePath, srcResp)
+    await uploadArtifactToS3(zipFilePath, srcResp, CodeWhispererConstants.FeatureUseCase.CODE_SCAN)
     getLogger().verbose(`Complete uploading src context.`)
     const artifactMap: ArtifactMap = {
         SourceCode: srcResp.uploadId,

diff --git a/packages/core/src/codewhisperer/service/securityScanHandler.ts b/packages/core/src/codewhisperer/service/securityScanHandler.ts
@@ -43,6 +43,8 @@ import { getTelemetryReasonDesc } from '../../shared/errors'
 import { CodeWhispererSettings } from '../util/codewhispererSettings'
 import { detectCommentAboveLine } from '../../shared/utilities/commentUtils'
 import { runtimeLanguageContext } from '../util/runtimeLanguageContext'
+import { FeatureUseCase } from '../models/constants'
+import { UploadTestArtifactToS3Error } from '../../amazonqTest/error'
 
 export async function listScanResults(
     client: DefaultCodeWhispererClient,
@@ -287,7 +289,7 @@ export async function getPresignedUrlAndUpload(
     logger.verbose(`CreateUploadUrlRequest request id: ${srcResp.$response.requestId}`)
     logger.verbose(`Complete Getting presigned Url for uploading src context.`)
     logger.verbose(`Uploading src context...`)
-    await uploadArtifactToS3(zipMetadata.zipFilePath, srcResp, scope)
+    await uploadArtifactToS3(zipMetadata.zipFilePath, srcResp, FeatureUseCase.CODE_SCAN, scope)
     logger.verbose(`Complete uploading src context.`)
     const artifactMap: ArtifactMap = {
         SourceCode: srcResp.uploadId,
@@ -343,6 +345,7 @@ export function throwIfCancelled(scope: CodeWhispererConstants.CodeAnalysisScope
 export async function uploadArtifactToS3(
     fileName: string,
     resp: CreateUploadUrlResponse,
+    featureUseCase: FeatureUseCase,
     scope?: CodeWhispererConstants.CodeAnalysisScope
 ) {
     const logger = getLoggerForScope(scope)
@@ -365,14 +368,23 @@ export async function uploadArtifactToS3(
         }).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 security scans. For more information, see the Amazon Q documentation or contact your network or organization administrator.`
+            `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 errorMessage = getTelemetryReasonDesc(error)?.includes(`"PUT" request failed with code "403"`)
-            ? `"PUT" request failed with code "403"`
-            : (getTelemetryReasonDesc(error) ?? 'Security scan failed.')
-
-        throw new UploadArtifactToS3Error(errorMessage)
+        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
+        }
+        throw isCodeScan ? new UploadArtifactToS3Error(errorMessage) : new UploadTestArtifactToS3Error(errorMessage)
     }
 }