fix(amazonq): Error handling and telemetry for Unit test generation.

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererCodeTestSession.kt

@@ -34,9 +34,6 @@ class CodeWhispererCodeTestSession(val sessionContext: CodeTestSessionContext) {
      * Run UTG sessions are follow steps:
      * 1. Zipping project
      * 2. Creating Upload url & Upload to S3 bucket
-     * 3. StartTestGeneration API -> Get JobId
-     * 4. GetTestGeneration API
-     * 5. ExportResultsArchieve API
      */
     suspend fun run(codeTestChatHelper: CodeTestChatHelper, previousIterationContext: PreviousUTGIterationContext?): CodeTestResponseContext {
         try {
@@ -94,7 +91,8 @@ class CodeWhispererCodeTestSession(val sessionContext: CodeTestSessionContext) {
                     sourceZip,
                     "SourceCode",
                     CodeWhispererConstants.UploadTaskType.UTG,
-                    taskName
+                    taskName,
+                    CodeWhispererConstants.FeatureName.TEST_GENERATION
                 )
 
             sourceZipUploadResponse.uploadId()
@@ -113,11 +111,10 @@ class CodeWhispererCodeTestSession(val sessionContext: CodeTestSessionContext) {
                 sourceZipUploadResponse,
                 testSummaryMessageId
             )
-            // TODO send telemetry for upload duration
 
             return codeTestResponseContext
         } catch (e: Exception) {
-            LOG.debug(e) { "Error when creating tests for the current file" }
+            LOG.debug(e) { "Error while creating zip and uploading to S3" }
             throw e
         }
     }

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt

@@ -16,7 +16,6 @@
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.launch
-import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException
 import software.amazon.awssdk.services.codewhispererruntime.model.GetTestGenerationResponse
 import software.amazon.awssdk.services.codewhispererruntime.model.Range
 import software.amazon.awssdk.services.codewhispererruntime.model.StartTestGenerationResponse
@@ -38,9 +37,14 @@
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.Session
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.utils.combineBuildAndExecuteLogFiles
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTotalLatency
+import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.CodeTestException
+import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.codeTestServerException
 import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.sessionconfig.CodeTestSessionConfig
+import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.testGenStoppedError
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
+import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.promptReAuth
+import software.aws.toolkits.jetbrains.services.codewhisperer.util.getTelemetryErrorMessage
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl
 import software.aws.toolkits.jetbrains.services.cwc.messages.ChatMessageType
 import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
@@ -67,7 +71,7 @@
 
     private fun throwIfCancelled(session: Session) {
         if (!session.isGeneratingTests) {
-            error(message("testgen.message.cancelled"))
+            throw testGenStoppedError()
         }
     }
 
@@ -104,24 +108,37 @@
         }
 
         // 2nd API call: StartTestGeneration
-        val startTestGenerationResponse = startTestGeneration(
-            uploadId = createUploadUrlResponse.uploadId(),
-            targetCode = listOf(
-                TargetCode.builder()
-                    .relativeTargetPath(codeTestResponseContext.currentFileRelativePath.toString())
-                    .targetLineRangeList(
-                        if (selectionRange != null) {
-                            listOf(
-                                selectionRange
-                            )
-                        } else {
-                            emptyList()
-                        }
-                    )
-                    .build()
-            ),
-            userInput = prompt
-        )
+        val startTestGenerationResponse = try {
+            startTestGeneration(
+                uploadId = createUploadUrlResponse.uploadId(),
+                targetCode = listOf(
+                    TargetCode.builder()
+                        .relativeTargetPath(codeTestResponseContext.currentFileRelativePath.toString())
+                        .targetLineRangeList(
+                            if (selectionRange != null) {
+                                listOf(
+                                    selectionRange
+                                )
+                            } else {
+                                emptyList()
+                            }
+                        )
+                        .build()
+                ),
+                userInput = prompt
+            )
+        } catch (e: Exception) {
+            LOG.error(e) { "Failed to create test generation job" }
+            // TODO: Not able to emit e.statusCode directly as statusCode is private property
+            // Cannot access 'statusCode': it is invisible (private in a supertype) in 'ThrottlingException'
+            val errorMessage = getTelemetryErrorMessage(e, CodeWhispererConstants.FeatureName.TEST_GENERATION)
+            throw codeTestServerException(
+                "CreateTestJobError: $errorMessage",
+                "400",
+                "CreateTestJobError",
+                message("testgen.error.generic_technical_error_message")
+            )
+        }
 
         val job = startTestGenerationResponse.testGenerationJob()
         session.startTestGenerationRequestId = startTestGenerationResponse.responseMetadata().requestId()
@@ -173,7 +190,13 @@
                     }
                     // update test summary card
                 } else {
-                    throw Exception(message("testgen.message.failed"))
+                    // If job status is Completed and has no ShortAnswer then there might be some issue in the backend.
+                    throw codeTestServerException(
+                        "TestGenFailedError: " + message("testgen.message.failed"),
+                        "500",
+                        "TestGenFailedError",
+                        message("testgen.error.generic_technical_error_message")
+                    )
                 }
             } else if (status == TestGenerationJobStatus.FAILED) {
                 LOG.debug {
@@ -184,11 +207,17 @@
                     shortAnswer = parseShortAnswerString(testGenerationResponse.testGenerationJob().shortAnswer())
                     if (shortAnswer.stopIteration == "true") {
                         throw Exception("${shortAnswer.planSummary}")
+                        throw codeTestServerException("TestGenFailedError: ${shortAnswer.planSummary}", "400", "TestGenFailedError", shortAnswer.planSummary)
                     }
                 }
 
-                // TODO: Modify text according to FnF
-                throw Exception(message("testgen.message.failed"))
+                // If job status is Failed and has no ShortAnswer then there might be some issue in the backend.
+                throw codeTestServerException(
+                    "TestGenFailedError: " + message("testgen.message.failed"),
+                    "500",
+                    "TestGenFailedError",
+                    message("testgen.error.generic_technical_error_message")
+                )
             } else {
                 // In progress
                 LOG.debug {
@@ -200,7 +229,7 @@
                 if (previousIterationContext == null && testGenerationResponse.testGenerationJob().shortAnswer() != null) {
                     shortAnswer = parseShortAnswerString(testGenerationResponse.testGenerationJob().shortAnswer())
                     if (shortAnswer.stopIteration == "true") {
-                        throw Exception("${shortAnswer.planSummary}")
+                        throw codeTestServerException("TestGenFailedError: ${shortAnswer.planSummary}", "400", "TestGenFailedError", shortAnswer.planSummary)
                     }
                     codeTestChatHelper.updateAnswer(
                         CodeTestChatMessageContent(
@@ -232,6 +261,12 @@
             },
             { e ->
                 LOG.error(e) { "ExportResultArchive failed: ${e.message}" }
+                throw codeTestServerException(
+                    "ExportResultsArchiveError: ${e.message}",
+                    "500",
+                    "ExportResultsArchiveError",
+                    message("testgen.error.generic_technical_error_message")
+                )
             },
             { startTime ->
                 LOG.info { "ExportResultArchive latency: ${calculateTotalLatency(startTime, Instant.now())}" }
@@ -493,14 +528,19 @@
                 launchTestGenFlow(prompt, codeTestChatHelper, previousIterationContext, selectionRange)
             } catch (e: Exception) {
                 // Add an answer for displaying error message
-                var errorMessage = e.message
-                if (e is JsonParseException) {
-                    errorMessage = message("testgen.error.generic_error_message")
+                val errorMessage = when {
+                    e is CodeTestException && e.statusCode == "400" &&
+                        e.message?.startsWith(
+                            "CreateTestJobError: " +
+                                "Maximum com.amazon.aws.codewhisperer.runtime.StartTestGeneration reached for this month."
+                        ) == true ->
+                        message("testgen.error.maximum_generations_reach")
+
+                    e is CodeTestException -> e.uiMessage
+                    e is JsonParseException -> message("testgen.error.generic_technical_error_message")
+                    else -> e.message
                 }
 
-                if (e is CodeWhispererRuntimeException) {
-                    errorMessage = message("testgen.error.maximum_generations_reach")
-                }
                 codeTestChatHelper.addAnswer(
                     CodeTestChatMessageContent(
                         message = errorMessage,
@@ -516,11 +556,12 @@
                     credentialStartUrl = getStartUrl(project),
                     jobGroup = session.testGenerationJobGroupName,
                     jobId = session.testGenerationJob,
                     result = if (e.message == message("testgen.message.cancelled")) MetricResult.Cancelled else MetricResult.Failed,
-                    reason = e.javaClass.name,
-                    reasonDesc = e.message,
+                    reason = (e as CodeTestException).code ?: "DefaultError",
+                    reasonDesc = if (e.message == message("testgen.message.cancelled")) "${e.code}: ${e.message}" else e.message,
+                    httpStatusCode = e.statusCode ?: "400",
                     perfClientLatency = (Instant.now().toEpochMilli() - session.startTimeOfTestGeneration),
                     isCodeBlockSelected = session.isCodeBlockSelected,
                     artifactsUploadDuration = session.artifactUploadDuration,
                     buildPayloadBytes = session.srcPayloadSize,
                     buildZipFileBytes = session.srcZipFileSize,

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt

@@ -291,7 +291,8 @@ class CodeTestChatController(
                     credentialStartUrl = getStartUrl(project),
                     result = MetricResult.Succeeded,
                     perfClientLatency = (Instant.now().toEpochMilli() - session.startTimeOfTestGeneration),
-                    requestId = id
+                    requestId = id,
+                    httpStatusCode = "200"
                 )
             }
             session.isGeneratingTests = false

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/AmazonQCodeFixSession.kt

@@ -72,7 +72,8 @@ class AmazonQCodeFixSession(val project: Project) {
                     sourceZip,
                     "SourceCode",
                     CodeWhispererConstants.UploadTaskType.CODE_FIX,
-                    codeFixName
+                    codeFixName,
+                    CodeWhispererConstants.FeatureName.CODE_REVIEW
                 )
 
             /**
@@ -149,7 +150,7 @@ class AmazonQCodeFixSession(val project: Project) {
         )
     } catch (e: Exception) {
         LOG.debug { "Create Upload URL failed: ${e.message}" }
-        val errorMessage = getTelemetryErrorMessage(e)
+        val errorMessage = getTelemetryErrorMessage(e, featureUseCase = CodeWhispererConstants.FeatureName.CODE_REVIEW)
         throw codeScanServerException("CreateUploadUrlException: $errorMessage")
     }
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanSession.kt

@@ -117,7 +117,8 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
                     sourceZip,
                     "SourceCode",
                     taskType,
-                    codeScanName
+                    codeScanName,
+                    CodeWhispererConstants.FeatureName.CODE_REVIEW
                 )
             if (isProjectScope()) {
                 LOG.debug {
@@ -272,7 +273,7 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
             )
         } catch (e: Exception) {
             LOG.debug { "Creating code review failed: ${e.message}" }
-            val errorMessage = getTelemetryErrorMessage(e)
+            val errorMessage = getTelemetryErrorMessage(e, featureUseCase = CodeWhispererConstants.FeatureName.CODE_REVIEW)
             throw codeScanServerException(errorMessage)
         }
     }
@@ -285,7 +286,7 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
         )
     } catch (e: Exception) {
         LOG.debug { "Getting code review failed: ${e.message}" }
-        val errorMessage = getTelemetryErrorMessage(e)
+        val errorMessage = getTelemetryErrorMessage(e, featureUseCase = CodeWhispererConstants.FeatureName.CODE_REVIEW)
         throw codeScanServerException("GetCodeReviewException: $errorMessage")
     }
 
@@ -299,7 +300,7 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
         )
     } catch (e: Exception) {
         LOG.debug { "Listing code review failed: ${e.message}" }
-        val errorMessage = getTelemetryErrorMessage(e)
+        val errorMessage = getTelemetryErrorMessage(e, featureUseCase = CodeWhispererConstants.FeatureName.CODE_REVIEW)
         throw codeScanServerException("ListCodeReviewFindingsException: $errorMessage")
     }
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codetest/CodeTestException.kt

@@ -0,0 +1,48 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.codewhisperer.codetest
+
+import software.aws.toolkits.resources.message
+
+open class CodeTestException(
+    override val message: String?,
+    val statusCode: String? = "400",
+    val code: String? = "DefaultError",
+    val uiMessage: String? = message(
+        "testgen.error.generic_error_message"
+    ),
+) : RuntimeException()
+
+open class CodeTestServerException(
+    override val message: String?,
+    val statusCode: String? = "400",
+    val code: String? = "DefaultError",
+    val uiMessage: String? = message(
+        "testgen.error.generic_technical_error_message"
+    ),
+) : RuntimeException()
+
+internal fun noFileOpenError(): Nothing =
+    throw CodeTestException(message("codewhisperer.codescan.no_file_open"), "400", "ProjectZipError")
+
+internal fun fileTooLarge(): Nothing =
+    throw CodeTestException(message("codewhisperer.codescan.file_too_large_telemetry"), "400", "ProjectZipError")
+
+internal fun cannotFindFile(errorMessage: String, filepath: String): Nothing =
+    error(message("codewhisperer.codescan.file_not_found", filepath, errorMessage))
+
+internal fun cannotFindValidFile(errorMessage: String): Nothing =
+    throw CodeTestException(errorMessage, "400", "ProjectZipError")
+
+internal fun cannotFindBuildArtifacts(errorMessage: String): Nothing =
+    throw CodeTestException(errorMessage, "400", "ProjectZipError")
+
+internal fun invalidSourceZipError(): Nothing =
+    throw CodeTestException(message("codewhisperer.codescan.invalid_source_zip_telemetry"), "400", "InvalidSourceZipError")
+
+fun codeTestServerException(errorMessage: String, statusCode: String?, code: String?, uiMessage: String?): Nothing =
+    throw CodeTestServerException(errorMessage, statusCode, code, uiMessage)
+
+fun testGenStoppedError(): Nothing =
+    throw CodeTestException(message("testgen.message.cancelled"), "400", "TestGenCancelled", message("testgen.message.cancelled"))

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codetest/sessionconfig/CodeTestSessionConfig.kt

@@ -19,13 +19,14 @@ import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.putNextEntry
 import software.aws.toolkits.jetbrains.services.amazonq.FeatureDevSessionContext
-import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.cannotFindBuildArtifacts
-import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.cannotFindFile
-import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.fileTooLarge
-import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.noFileOpenError
 import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.Payload
 import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.PayloadContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.PayloadMetadata
+import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.cannotFindBuildArtifacts
+import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.cannotFindFile
+import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.cannotFindValidFile
+import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.fileTooLarge
+import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.noFileOpenError
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererUnknownLanguage
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmingLanguage
@@ -94,7 +95,7 @@ class CodeTestSessionConfig(
                 else -> e.message
             }
             LOG.debug { "Error creating payload metadata: $errorMessage" }
-            throw cannotFindBuildArtifacts(errorMessage ?: message("codewhisperer.codescan.run_scan_error_telemetry"))
+            throw cannotFindBuildArtifacts(errorMessage ?: message("testgen.message.failed"))
         }
 
         // Copy all the included source files to the source zip
@@ -218,7 +219,7 @@ class CodeTestSessionConfig(
 
         if (maxCountLanguage == null) {
             programmingLanguage = CodeWhispererUnknownLanguage.INSTANCE
-            throw RuntimeException("Amazon Q: doesn't contain valid files to generate tests")
+            throw cannotFindValidFile("Amazon Q: doesn't contain valid files to generate tests")
         }
         programmingLanguage = maxCountLanguage
         return PayloadMetadata(files, currentTotalFileSize, currentTotalLines, maxCountLanguage.toTelemetryType())

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt

@@ -127,6 +127,11 @@ object CodeWhispererConstants {
         PROJECT("PROJECT"),
     }
 
+    enum class FeatureName(val value: String) {
+        TEST_GENERATION("TEST_GENERATION"),
+        CODE_REVIEW("CODE_REVIEW"),
+    }
+
     enum class UploadTaskType(val value: String) {
         SCAN_FILE("SCAN_FILE"),
         SCAN_PROJECT("SCAN_PROJECT"),