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,12 +16,12 @@
 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
 import software.amazon.awssdk.services.codewhispererruntime.model.TargetCode
 import software.amazon.awssdk.services.codewhispererruntime.model.TestGenerationJobStatus
+import software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingException
 import software.amazon.awssdk.services.codewhispererstreaming.model.ExportContext
 import software.amazon.awssdk.services.codewhispererstreaming.model.ExportIntent
 import software.aws.toolkits.core.utils.debug
@@ -38,9 +38,13 @@
 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.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"))
+            testGenStoppedError()
         }
     }
 
@@ -104,24 +108,39 @@
         }
 
         // 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) {
+            val statusCode = when {
+                e is ThrottlingException -> e.statusCode()
+                else -> 500
+            }
+            LOG.error(e) { "Unexpected error while creating test generation job" }
+            val errorMessage = getTelemetryErrorMessage(e, CodeWhispererConstants.FeatureName.TEST_GENERATION)
+            throw CodeTestException(
+                "CreateTestJobError: $errorMessage",
+                statusCode,
+                "CreateTestJobError",
+                message("testgen.error.generic_technical_error_message")
+            )
+        }
 
         val job = startTestGenerationResponse.testGenerationJob()
         session.startTestGenerationRequestId = startTestGenerationResponse.responseMetadata().requestId()
@@ -173,7 +192,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 CodeTestException(
+                        "TestGenFailedError: " + message("testgen.message.failed"),
+                        500,
+                        "TestGenFailedError",
+                        message("testgen.error.generic_technical_error_message")
+                    )
                 }
             } else if (status == TestGenerationJobStatus.FAILED) {
                 LOG.debug {
@@ -183,12 +208,17 @@
                 if (testGenerationResponse.testGenerationJob().shortAnswer() != null) {
                     shortAnswer = parseShortAnswerString(testGenerationResponse.testGenerationJob().shortAnswer())
                     if (shortAnswer.stopIteration == "true") {
-                        throw Exception("${shortAnswer.planSummary}")
+                        throw CodeTestException("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 CodeTestException(
+                    "TestGenFailedError: " + message("testgen.message.failed"),
+                    500,
+                    "TestGenFailedError",
+                    message("testgen.error.generic_technical_error_message")
+                )
             } else {
                 // In progress
                 LOG.debug {
@@ -200,7 +230,7 @@
                 if (previousIterationContext == null && testGenerationResponse.testGenerationJob().shortAnswer() != null) {
                     shortAnswer = parseShortAnswerString(testGenerationResponse.testGenerationJob().shortAnswer())
                     if (shortAnswer.stopIteration == "true") {
-                        throw Exception("${shortAnswer.planSummary}")
+                        throw CodeTestException("TestGenFailedError: ${shortAnswer.planSummary}", 400, "TestGenFailedError", shortAnswer.planSummary)
                     }
                     codeTestChatHelper.updateAnswer(
                         CodeTestChatMessageContent(
@@ -232,6 +262,12 @@
             },
             { e ->
                 LOG.error(e) { "ExportResultArchive failed: ${e.message}" }
+                throw CodeTestException(
+                    "ExportResultsArchiveError: ${e.message}",
+                    500,
+                    "ExportResultsArchiveError",
+                    message("testgen.error.generic_technical_error_message")
+                )
             },
             { startTime ->
                 LOG.info { "ExportResultArchive latency: ${calculateTotalLatency(startTime, Instant.now())}" }
@@ -493,14 +529,16 @@
                 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") == true ->
+                        message("testgen.error.maximum_generations_reach")
+
+                    e is CodeTestException -> e.uiMessage
+                    e is JsonParseException -> message("testgen.error.generic_technical_error_message")
+                    else -> message("testgen.error.generic_error_message")
                 }
 
-                if (e is CodeWhispererRuntimeException) {
-                    errorMessage = message("testgen.error.maximum_generations_reach")
-                }
                 codeTestChatHelper.addAnswer(
                     CodeTestChatMessageContent(
                         message = errorMessage,
@@ -517,8 +555,9 @@
                     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.toString(),
                     perfClientLatency = (Instant.now().toEpochMilli() - session.startTimeOfTestGeneration),
                     isCodeBlockSelected = session.isCodeBlockSelected,
                     artifactsUploadDuration = session.artifactUploadDuration,

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
@@ -442,7 +443,6 @@ class CodeTestChatController(
         var charDifference = 0
         var generatedFileContent = ""
         var selectedFileContent = ""
-        var latencyOfTestGeneration = 0.0
 
         when (message.actionID) {
             "utg_view_diff" -> {
@@ -486,7 +486,7 @@ class CodeTestChatController(
 
                     session.linesOfCodeGenerated = lineDifference.coerceAtLeast(0)
                     session.charsOfCodeGenerated = charDifference.coerceAtLeast(0)
-                    latencyOfTestGeneration = (Instant.now().toEpochMilli() - session.startTimeOfTestGeneration)
+                    session.latencyOfTestGeneration = (Instant.now().toEpochMilli() - session.startTimeOfTestGeneration)
                     UiTelemetry.click(null as Project?, "unitTestGeneration_viewDiff")
 
                     val buttonList = mutableListOf<Button>()
@@ -599,12 +599,13 @@ class CodeTestChatController(
                     acceptedCharactersCount = session.charsOfCodeGenerated?.toLong(),
                     generatedCharactersCount = session.charsOfCodeGenerated?.toLong(),
                     result = MetricResult.Succeeded,
-                    perfClientLatency = latencyOfTestGeneration,
+                    perfClientLatency = session.latencyOfTestGeneration,
                     isCodeBlockSelected = session.isCodeBlockSelected,
                     artifactsUploadDuration = session.artifactUploadDuration,
                     buildPayloadBytes = session.srcPayloadSize,
                     buildZipFileBytes = session.srcZipFileSize,
-                    requestId = session.startTestGenerationRequestId
+                    requestId = session.startTestGenerationRequestId,
+                    httpStatusCode = "200"
                 )
                 codeTestChatHelper.addAnswer(
                     CodeTestChatMessageContent(
@@ -793,12 +794,13 @@ class CodeTestChatController(
                     acceptedCharactersCount = 0,
                     generatedCharactersCount = session.charsOfCodeGenerated?.toLong(),
                     result = MetricResult.Succeeded,
-                    perfClientLatency = latencyOfTestGeneration,
+                    perfClientLatency = session.latencyOfTestGeneration,
                     isCodeBlockSelected = session.isCodeBlockSelected,
                     artifactsUploadDuration = session.artifactUploadDuration,
                     buildPayloadBytes = session.srcPayloadSize,
                     buildZipFileBytes = session.srcZipFileSize,
-                    requestId = session.startTestGenerationRequestId
+                    requestId = session.startTestGenerationRequestId,
+                    httpStatusCode = "200"
                 )
                 sessionCleanUp(message.tabId)
             }

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

@@ -28,6 +28,7 @@ data class Session(val tabId: String) {
     var linesOfCodeGenerated: Int? = null
     var charsOfCodeGenerated: Int? = null
     var startTimeOfTestGeneration: Double = 0.0
+    var latencyOfTestGeneration: Double = 0.0
     var isCodeBlockSelected: Boolean = false
     var srcPayloadSize: Long = 0
     var srcZipFileSize: Long = 0

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,36 @@
+// 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: Int = 400,
+    val code: String? = "DefaultError",
+    val uiMessage: String? = message(
+        "testgen.error.generic_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 testGenStoppedError(): Nothing =
+    throw CodeTestException(message("testgen.message.cancelled"), 400, "TestGenCancelled", message("testgen.message.cancelled"))