Skip to content

Commit 4053492

Browse files
fix(amazonq): Error handling and telemetry for Unit test generation. (#5192)
* Error handling for UTG Co-authored-by: invictus <[email protected]>
1 parent 8093096 commit 4053492

File tree

14 files changed

+201
-85
lines changed

14 files changed

+201
-85
lines changed

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

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,6 @@ class CodeWhispererCodeTestSession(val sessionContext: CodeTestSessionContext) {
3434
* Run UTG sessions are follow steps:
3535
* 1. Zipping project
3636
* 2. Creating Upload url & Upload to S3 bucket
37-
* 3. StartTestGeneration API -> Get JobId
38-
* 4. GetTestGeneration API
39-
* 5. ExportResultsArchieve API
4037
*/
4138
suspend fun run(codeTestChatHelper: CodeTestChatHelper, previousIterationContext: PreviousUTGIterationContext?): CodeTestResponseContext {
4239
try {
@@ -94,7 +91,8 @@ class CodeWhispererCodeTestSession(val sessionContext: CodeTestSessionContext) {
9491
sourceZip,
9592
"SourceCode",
9693
CodeWhispererConstants.UploadTaskType.UTG,
97-
taskName
94+
taskName,
95+
CodeWhispererConstants.FeatureName.TEST_GENERATION
9896
)
9997

10098
sourceZipUploadResponse.uploadId()
@@ -113,11 +111,10 @@ class CodeWhispererCodeTestSession(val sessionContext: CodeTestSessionContext) {
113111
sourceZipUploadResponse,
114112
testSummaryMessageId
115113
)
116-
// TODO send telemetry for upload duration
117114

118115
return codeTestResponseContext
119116
} catch (e: Exception) {
120-
LOG.debug(e) { "Error when creating tests for the current file" }
117+
LOG.debug(e) { "Error while creating zip and uploading to S3" }
121118
throw e
122119
}
123120
}

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

Lines changed: 67 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ import kotlinx.coroutines.CoroutineScope
1616
import kotlinx.coroutines.Job
1717
import kotlinx.coroutines.delay
1818
import kotlinx.coroutines.launch
19-
import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException
19+
import software.amazon.awssdk.core.exception.SdkServiceException
2020
import software.amazon.awssdk.services.codewhispererruntime.model.GetTestGenerationResponse
2121
import software.amazon.awssdk.services.codewhispererruntime.model.Range
2222
import software.amazon.awssdk.services.codewhispererruntime.model.StartTestGenerationResponse
@@ -38,9 +38,13 @@ import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.BuildAnd
3838
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.Session
3939
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.utils.combineBuildAndExecuteLogFiles
4040
import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTotalLatency
41+
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.CodeTestException
4142
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.sessionconfig.CodeTestSessionConfig
43+
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.testGenStoppedError
4244
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
45+
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
4346
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.promptReAuth
47+
import software.aws.toolkits.jetbrains.services.codewhisperer.util.getTelemetryErrorMessage
4448
import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl
4549
import software.aws.toolkits.jetbrains.services.cwc.messages.ChatMessageType
4650
import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
@@ -67,7 +71,7 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
6771

6872
private fun throwIfCancelled(session: Session) {
6973
if (!session.isGeneratingTests) {
70-
error(message("testgen.message.cancelled"))
74+
testGenStoppedError()
7175
}
7276
}
7377

@@ -104,24 +108,38 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
104108
}
105109

106110
// 2nd API call: StartTestGeneration
107-
val startTestGenerationResponse = startTestGeneration(
108-
uploadId = createUploadUrlResponse.uploadId(),
109-
targetCode = listOf(
110-
TargetCode.builder()
111-
.relativeTargetPath(codeTestResponseContext.currentFileRelativePath.toString())
112-
.targetLineRangeList(
113-
if (selectionRange != null) {
114-
listOf(
115-
selectionRange
116-
)
117-
} else {
118-
emptyList()
119-
}
120-
)
121-
.build()
122-
),
123-
userInput = prompt
124-
)
111+
val startTestGenerationResponse = try {
112+
startTestGeneration(
113+
uploadId = createUploadUrlResponse.uploadId(),
114+
targetCode = listOf(
115+
TargetCode.builder()
116+
.relativeTargetPath(codeTestResponseContext.currentFileRelativePath.toString())
117+
.targetLineRangeList(
118+
if (selectionRange != null) {
119+
listOf(
120+
selectionRange
121+
)
122+
} else {
123+
emptyList()
124+
}
125+
)
126+
.build()
127+
),
128+
userInput = prompt
129+
)
130+
} catch (e: Exception) {
131+
val statusCode = when {
132+
e is SdkServiceException -> e.statusCode()
133+
else -> 400
134+
}
135+
LOG.error(e) { "Unexpected error while creating test generation job" }
136+
val errorMessage = getTelemetryErrorMessage(e, CodeWhispererConstants.FeatureName.TEST_GENERATION)
137+
throw CodeTestException(
138+
"CreateTestJobError: $errorMessage",
139+
"CreateTestJobError",
140+
message("testgen.error.generic_technical_error_message")
141+
)
142+
}
125143

126144
val job = startTestGenerationResponse.testGenerationJob()
127145
session.startTestGenerationRequestId = startTestGenerationResponse.responseMetadata().requestId()
@@ -173,7 +191,12 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
173191
}
174192
// update test summary card
175193
} else {
176-
throw Exception(message("testgen.message.failed"))
194+
// If job status is Completed and has no ShortAnswer then there might be some issue in the backend.
195+
throw CodeTestException(
196+
"TestGenFailedError: " + message("testgen.message.failed"),
197+
"TestGenFailedError",
198+
message("testgen.error.generic_technical_error_message")
199+
)
177200
}
178201
} else if (status == TestGenerationJobStatus.FAILED) {
179202
LOG.debug {
@@ -183,12 +206,16 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
183206
if (testGenerationResponse.testGenerationJob().shortAnswer() != null) {
184207
shortAnswer = parseShortAnswerString(testGenerationResponse.testGenerationJob().shortAnswer())
185208
if (shortAnswer.stopIteration == "true") {
186-
throw Exception("${shortAnswer.planSummary}")
209+
throw CodeTestException("TestGenFailedError: ${shortAnswer.planSummary}", "TestGenFailedError", shortAnswer.planSummary)
187210
}
188211
}
189212

190-
// TODO: Modify text according to FnF
191-
throw Exception(message("testgen.message.failed"))
213+
// If job status is Failed and has no ShortAnswer then there might be some issue in the backend.
214+
throw CodeTestException(
215+
"TestGenFailedError: " + message("testgen.message.failed"),
216+
"TestGenFailedError",
217+
message("testgen.error.generic_technical_error_message")
218+
)
192219
} else {
193220
// In progress
194221
LOG.debug {
@@ -200,7 +227,7 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
200227
if (previousIterationContext == null && testGenerationResponse.testGenerationJob().shortAnswer() != null) {
201228
shortAnswer = parseShortAnswerString(testGenerationResponse.testGenerationJob().shortAnswer())
202229
if (shortAnswer.stopIteration == "true") {
203-
throw Exception("${shortAnswer.planSummary}")
230+
throw CodeTestException("TestGenFailedError: ${shortAnswer.planSummary}", "TestGenFailedError", shortAnswer.planSummary)
204231
}
205232
codeTestChatHelper.updateAnswer(
206233
CodeTestChatMessageContent(
@@ -232,6 +259,11 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
232259
},
233260
{ e ->
234261
LOG.error(e) { "ExportResultArchive failed: ${e.message}" }
262+
throw CodeTestException(
263+
"ExportResultsArchiveError: ${e.message}",
264+
"ExportResultsArchiveError",
265+
message("testgen.error.generic_technical_error_message")
266+
)
235267
},
236268
{ startTime ->
237269
LOG.info { "ExportResultArchive latency: ${calculateTotalLatency(startTime, Instant.now())}" }
@@ -493,14 +525,16 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
493525
launchTestGenFlow(prompt, codeTestChatHelper, previousIterationContext, selectionRange)
494526
} catch (e: Exception) {
495527
// Add an answer for displaying error message
496-
var errorMessage = e.message
497-
if (e is JsonParseException) {
498-
errorMessage = message("testgen.error.generic_error_message")
528+
val errorMessage = when {
529+
e is CodeTestException &&
530+
e.message?.startsWith("CreateTestJobError: Maximum") == true ->
531+
message("testgen.error.maximum_generations_reach")
532+
533+
e is CodeTestException -> e.uiMessage
534+
e is JsonParseException -> message("testgen.error.generic_technical_error_message")
535+
else -> message("testgen.error.generic_error_message")
499536
}
500537

501-
if (e is CodeWhispererRuntimeException) {
502-
errorMessage = message("testgen.error.maximum_generations_reach")
503-
}
504538
codeTestChatHelper.addAnswer(
505539
CodeTestChatMessageContent(
506540
message = errorMessage,
@@ -518,8 +552,8 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
518552
jobGroup = session.testGenerationJobGroupName,
519553
jobId = session.testGenerationJob,
520554
result = if (e.message == message("testgen.message.cancelled")) MetricResult.Cancelled else MetricResult.Failed,
521-
reason = e.javaClass.name,
522-
reasonDesc = e.message,
555+
reason = (e as CodeTestException).code ?: "DefaultError",
556+
reasonDesc = if (e.message == message("testgen.message.cancelled")) "${e.code}: ${e.message}" else e.message,
523557
perfClientLatency = (Instant.now().toEpochMilli() - session.startTimeOfTestGeneration),
524558
isCodeBlockSelected = session.isCodeBlockSelected,
525559
artifactsUploadDuration = session.artifactUploadDuration,

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -452,7 +452,6 @@ class CodeTestChatController(
452452
var charDifference = 0
453453
var generatedFileContent = ""
454454
var selectedFileContent = ""
455-
var latencyOfTestGeneration = 0.0
456455

457456
when (message.actionID) {
458457
"utg_view_diff" -> {
@@ -496,7 +495,7 @@ class CodeTestChatController(
496495

497496
session.linesOfCodeGenerated = lineDifference.coerceAtLeast(0)
498497
session.charsOfCodeGenerated = charDifference.coerceAtLeast(0)
499-
latencyOfTestGeneration = (Instant.now().toEpochMilli() - session.startTimeOfTestGeneration)
498+
session.latencyOfTestGeneration = (Instant.now().toEpochMilli() - session.startTimeOfTestGeneration)
500499
UiTelemetry.click(null as Project?, "unitTestGeneration_viewDiff")
501500

502501
val buttonList = mutableListOf<Button>()
@@ -611,7 +610,7 @@ class CodeTestChatController(
611610
acceptedCharactersCount = session.charsOfCodeGenerated?.toLong(),
612611
generatedCharactersCount = session.charsOfCodeGenerated?.toLong(),
613612
result = MetricResult.Succeeded,
614-
perfClientLatency = latencyOfTestGeneration,
613+
perfClientLatency = session.latencyOfTestGeneration,
615614
isCodeBlockSelected = session.isCodeBlockSelected,
616615
artifactsUploadDuration = session.artifactUploadDuration,
617616
buildPayloadBytes = session.srcPayloadSize,
@@ -807,7 +806,7 @@ class CodeTestChatController(
807806
acceptedCharactersCount = 0,
808807
generatedCharactersCount = session.charsOfCodeGenerated?.toLong(),
809808
result = MetricResult.Succeeded,
810-
perfClientLatency = latencyOfTestGeneration,
809+
perfClientLatency = session.latencyOfTestGeneration,
811810
isCodeBlockSelected = session.isCodeBlockSelected,
812811
artifactsUploadDuration = session.artifactUploadDuration,
813812
buildPayloadBytes = session.srcPayloadSize,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ data class Session(val tabId: String) {
2828
var linesOfCodeGenerated: Int? = null
2929
var charsOfCodeGenerated: Int? = null
3030
var startTimeOfTestGeneration: Double = 0.0
31+
var latencyOfTestGeneration: Double = 0.0
3132
var isCodeBlockSelected: Boolean = false
3233
var srcPayloadSize: Long = 0
3334
var srcZipFileSize: Long = 0

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,8 @@ class AmazonQCodeFixSession(val project: Project) {
7272
sourceZip,
7373
"SourceCode",
7474
CodeWhispererConstants.UploadTaskType.CODE_FIX,
75-
codeFixName
75+
codeFixName,
76+
CodeWhispererConstants.FeatureName.CODE_REVIEW
7677
)
7778

7879
/**
@@ -149,7 +150,7 @@ class AmazonQCodeFixSession(val project: Project) {
149150
)
150151
} catch (e: Exception) {
151152
LOG.debug { "Create Upload URL failed: ${e.message}" }
152-
val errorMessage = getTelemetryErrorMessage(e)
153+
val errorMessage = getTelemetryErrorMessage(e, featureUseCase = CodeWhispererConstants.FeatureName.CODE_REVIEW)
153154
throw codeScanServerException("CreateUploadUrlException: $errorMessage")
154155
}
155156

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,8 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
117117
sourceZip,
118118
"SourceCode",
119119
taskType,
120-
codeScanName
120+
codeScanName,
121+
CodeWhispererConstants.FeatureName.CODE_REVIEW
121122
)
122123
if (isProjectScope()) {
123124
LOG.debug {
@@ -272,7 +273,7 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
272273
)
273274
} catch (e: Exception) {
274275
LOG.debug { "Creating code review failed: ${e.message}" }
275-
val errorMessage = getTelemetryErrorMessage(e)
276+
val errorMessage = getTelemetryErrorMessage(e, featureUseCase = CodeWhispererConstants.FeatureName.CODE_REVIEW)
276277
throw codeScanServerException(errorMessage)
277278
}
278279
}
@@ -285,7 +286,7 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
285286
)
286287
} catch (e: Exception) {
287288
LOG.debug { "Getting code review failed: ${e.message}" }
288-
val errorMessage = getTelemetryErrorMessage(e)
289+
val errorMessage = getTelemetryErrorMessage(e, featureUseCase = CodeWhispererConstants.FeatureName.CODE_REVIEW)
289290
throw codeScanServerException("GetCodeReviewException: $errorMessage")
290291
}
291292

@@ -299,7 +300,7 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
299300
)
300301
} catch (e: Exception) {
301302
LOG.debug { "Listing code review failed: ${e.message}" }
302-
val errorMessage = getTelemetryErrorMessage(e)
303+
val errorMessage = getTelemetryErrorMessage(e, featureUseCase = CodeWhispererConstants.FeatureName.CODE_REVIEW)
303304
throw codeScanServerException("ListCodeReviewFindingsException: $errorMessage")
304305
}
305306

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.codewhisperer.codetest
5+
6+
import software.aws.toolkits.resources.message
7+
8+
open class CodeTestException(
9+
override val message: String?,
10+
val code: String? = "DefaultError",
11+
val uiMessage: String? = message(
12+
"testgen.error.generic_error_message"
13+
),
14+
) : RuntimeException()
15+
16+
internal fun noFileOpenError(): Nothing =
17+
throw CodeTestException(message("codewhisperer.codescan.no_file_open"), "ProjectZipError")
18+
19+
internal fun fileTooLarge(): Nothing =
20+
throw CodeTestException(message("codewhisperer.codescan.file_too_large_telemetry"), "ProjectZipError")
21+
22+
internal fun cannotFindFile(errorMessage: String, filepath: String): Nothing =
23+
error(message("codewhisperer.codescan.file_not_found", filepath, errorMessage))
24+
25+
internal fun cannotFindValidFile(errorMessage: String): Nothing =
26+
throw CodeTestException(errorMessage, "ProjectZipError")
27+
28+
internal fun cannotFindBuildArtifacts(errorMessage: String): Nothing =
29+
throw CodeTestException(errorMessage, "ProjectZipError")
30+
31+
internal fun invalidSourceZipError(): Nothing =
32+
throw CodeTestException(message("codewhisperer.codescan.invalid_source_zip_telemetry"), "InvalidSourceZipError")
33+
34+
fun testGenStoppedError(): Nothing =
35+
throw CodeTestException(message("testgen.message.cancelled"), "TestGenCancelled", message("testgen.message.cancelled"))

0 commit comments

Comments
 (0)