Skip to content

Commit 9afffbd

Browse files
evanliu048evanliu048
authored
Merge branch 'main' into revisePrefetchSession
2 parents 5cf9b14 + 17dcd44 commit 9afffbd

File tree

14 files changed

+369
-141
lines changed

14 files changed

+369
-141
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,12 +75,12 @@ class CodeWhispererCodeTestSession(val sessionContext: CodeTestSessionContext) {
7575

7676
LOG.debug {
7777
"Total size of source payload in KB: ${payloadContext.srcPayloadSize * 1.0 / TOTAL_BYTES_IN_KB} \n" +
78-
"Total size of build payload in KB: ${(payloadContext.buildPayloadSize ?: 0) * 1.0 / TOTAL_BYTES_IN_KB} \n" +
7978
"Total size of source zip file in KB: ${payloadContext.srcZipFileSize * 1.0 / TOTAL_BYTES_IN_KB} \n" +
8079
"Total number of lines included: ${payloadContext.totalLines} \n" +
8180
"Total number of files included in payload: ${payloadContext.totalFiles} \n" +
8281
"Total time taken for creating payload: ${payloadContext.totalTimeInMilliseconds * 1.0 / TOTAL_MILLIS_IN_SECOND} seconds\n" +
83-
"Payload context language: ${payloadContext.language}"
82+
"Payload context language: ${payloadContext.language}" +
83+
"Payload exceeded the limit: ${payloadContext.payloadLimitCrossed}"
8484
}
8585

8686
// 2 & 3. CreateUploadURL and upload the context.

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +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.core.exception.SdkServiceException
1920
import software.amazon.awssdk.services.codewhispererruntime.model.GetTestGenerationResponse
2021
import software.amazon.awssdk.services.codewhispererruntime.model.Range
2122
import software.amazon.awssdk.services.codewhispererruntime.model.StartTestGenerationResponse
@@ -39,6 +40,7 @@ import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.Session
3940
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.utils.combineBuildAndExecuteLogFiles
4041
import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTotalLatency
4142
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.CodeTestException
43+
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.fileTooLarge
4244
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.sessionconfig.CodeTestSessionConfig
4345
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.testGenStoppedError
4446
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
@@ -98,6 +100,9 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
98100
session.srcZipFileSize = codeTestResponseContext.payloadContext.srcZipFileSize
99101
session.artifactUploadDuration = codeTestResponseContext.serviceInvocationContext.artifactsUploadDuration
100102
val path = codeTestResponseContext.currentFileRelativePath
103+
if (codeTestResponseContext.payloadContext.payloadLimitCrossed == true) {
104+
fileTooLarge()
105+
}
101106

102107
val createUploadUrlResponse = codeTestResponseContext.createUploadUrlResponse ?: return
103108
throwIfCancelled(session)
@@ -144,6 +149,11 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
144149
} catch (e: Exception) {
145150
LOG.error(e) { "Unexpected error while creating test generation job" }
146151
val errorMessage = getTelemetryErrorMessage(e, CodeWhispererConstants.FeatureName.TEST_GENERATION)
152+
153+
// Sending requestId to telemetry if there is Validation Exception
154+
if (e is SdkServiceException) {
155+
session.startTestGenerationRequestId = e.requestId()
156+
}
147157
throw CodeTestException(
148158
"CreateTestJobError: $errorMessage",
149159
"CreateTestJobError",

‎plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/clients/AmazonQStreamingClientTest.kt‎

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,15 @@ package software.aws.toolkits.jetbrains.services.amazonq.clients
55

66
import com.intellij.testFramework.RuleChain
77
import com.intellij.testFramework.replaceService
8+
import kotlinx.coroutines.runBlocking
89
import kotlinx.coroutines.test.runTest
10+
import org.assertj.core.api.Assertions.assertThat
911
import org.junit.Before
1012
import org.junit.Rule
1113
import org.junit.Test
1214
import org.mockito.kotlin.any
1315
import org.mockito.kotlin.argumentCaptor
16+
import org.mockito.kotlin.doAnswer
1417
import org.mockito.kotlin.doReturn
1518
import org.mockito.kotlin.mock
1619
import org.mockito.kotlin.stub
@@ -20,6 +23,7 @@ import software.amazon.awssdk.services.codewhispererstreaming.CodeWhispererStrea
2023
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportIntent
2124
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportResultArchiveRequest
2225
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportResultArchiveResponseHandler
26+
import software.amazon.awssdk.services.codewhispererstreaming.model.ValidationException
2327
import software.amazon.awssdk.services.ssooidc.SsoOidcClient
2428
import software.aws.toolkits.core.TokenConnectionSettings
2529
import software.aws.toolkits.core.utils.test.aString
@@ -81,4 +85,156 @@ class AmazonQStreamingClientTest : AmazonQTestBase() {
8185
verify(streamingBearerClient).exportResultArchive(requestCaptor.capture(), handlerCaptor.capture())
8286
}
8387
}
88+
89+
@Test
90+
fun `verify retry on ValidationException`(): Unit = runBlocking {
91+
var attemptCount = 0
92+
streamingBearerClient = mockClientManagerRule.create<CodeWhispererStreamingAsyncClient>().stub {
93+
on {
94+
exportResultArchive(any<ExportResultArchiveRequest>(), any<ExportResultArchiveResponseHandler>())
95+
} doAnswer {
96+
attemptCount++
97+
if (attemptCount <= 2) {
98+
CompletableFuture.runAsync {
99+
throw VALIDATION_EXCEPTION
100+
}
101+
} else {
102+
CompletableFuture.completedFuture(mock())
103+
}
104+
}
105+
}
106+
107+
amazonQStreamingClient.exportResultArchive("test-id", ExportIntent.TRANSFORMATION, null, {}, {})
108+
109+
assertThat(attemptCount).isEqualTo(3)
110+
}
111+
112+
@Test
113+
fun `verify retry gives up after max attempts`(): Unit = runBlocking {
114+
var attemptCount = 0
115+
streamingBearerClient = mockClientManagerRule.create<CodeWhispererStreamingAsyncClient>().stub {
116+
on {
117+
exportResultArchive(any<ExportResultArchiveRequest>(), any<ExportResultArchiveResponseHandler>())
118+
} doAnswer {
119+
attemptCount++
120+
CompletableFuture.runAsync {
121+
throw VALIDATION_EXCEPTION
122+
}
123+
}
124+
}
125+
126+
val thrown = catchCoroutineException {
127+
amazonQStreamingClient.exportResultArchive("test-id", ExportIntent.TRANSFORMATION, null, {}, {})
128+
}
129+
130+
assertThat(attemptCount).isEqualTo(3)
131+
assertThat(thrown)
132+
.isInstanceOf(ValidationException::class.java)
133+
.hasMessage("Resource validation failed")
134+
}
135+
136+
@Test
137+
fun `verify no retry on non-retryable exception`(): Unit = runBlocking {
138+
var attemptCount = 0
139+
140+
streamingBearerClient = mockClientManagerRule.create<CodeWhispererStreamingAsyncClient>().stub {
141+
on {
142+
exportResultArchive(any<ExportResultArchiveRequest>(), any<ExportResultArchiveResponseHandler>())
143+
} doAnswer {
144+
attemptCount++
145+
CompletableFuture.runAsync {
146+
throw IllegalArgumentException("Non-retryable error")
147+
}
148+
}
149+
}
150+
151+
val thrown = catchCoroutineException {
152+
amazonQStreamingClient.exportResultArchive("test-id", ExportIntent.TRANSFORMATION, null, {}, {})
153+
}
154+
155+
assertThat(attemptCount).isEqualTo(1)
156+
assertThat(thrown)
157+
.isInstanceOf(IllegalArgumentException::class.java)
158+
.hasMessage("Non-retryable error")
159+
}
160+
161+
@Test
162+
fun `verify backoff timing between retries`(): Unit = runBlocking {
163+
var lastAttemptTime = 0L
164+
var minBackoffObserved = Long.MAX_VALUE
165+
var maxBackoffObserved = 0L
166+
167+
streamingBearerClient = mockClientManagerRule.create<CodeWhispererStreamingAsyncClient>().stub {
168+
on {
169+
exportResultArchive(any<ExportResultArchiveRequest>(), any<ExportResultArchiveResponseHandler>())
170+
} doAnswer {
171+
val currentTime = System.currentTimeMillis()
172+
if (lastAttemptTime > 0) {
173+
val backoffTime = currentTime - lastAttemptTime
174+
minBackoffObserved = minOf(minBackoffObserved, backoffTime)
175+
maxBackoffObserved = maxOf(maxBackoffObserved, backoffTime)
176+
}
177+
lastAttemptTime = currentTime
178+
179+
CompletableFuture.runAsync {
180+
throw VALIDATION_EXCEPTION
181+
}
182+
}
183+
}
184+
185+
val thrown = catchCoroutineException {
186+
amazonQStreamingClient.exportResultArchive("test-id", ExportIntent.TRANSFORMATION, null, {}, {})
187+
}
188+
189+
assertThat(thrown)
190+
.isInstanceOf(ValidationException::class.java)
191+
.hasMessage("Resource validation failed")
192+
assertThat(minBackoffObserved).isGreaterThanOrEqualTo(100)
193+
assertThat(maxBackoffObserved).isLessThanOrEqualTo(10000)
194+
}
195+
196+
@Test
197+
fun `verify onError callback is called with final exception`(): Unit = runBlocking {
198+
var errorCaught: Exception? = null
199+
200+
streamingBearerClient = mockClientManagerRule.create<CodeWhispererStreamingAsyncClient>().stub {
201+
on {
202+
exportResultArchive(any<ExportResultArchiveRequest>(), any<ExportResultArchiveResponseHandler>())
203+
} doAnswer {
204+
CompletableFuture.runAsync {
205+
throw VALIDATION_EXCEPTION
206+
}
207+
}
208+
}
209+
210+
val thrown = catchCoroutineException {
211+
amazonQStreamingClient.exportResultArchive(
212+
"test-id",
213+
ExportIntent.TRANSFORMATION,
214+
null,
215+
{ errorCaught = it },
216+
{}
217+
)
218+
}
219+
220+
assertThat(thrown)
221+
.isInstanceOf(ValidationException::class.java)
222+
.hasMessage("Resource validation failed")
223+
assertThat(errorCaught).isEqualTo(VALIDATION_EXCEPTION)
224+
}
225+
226+
private suspend fun catchCoroutineException(block: suspend () -> Unit): Throwable {
227+
try {
228+
block()
229+
error("Expected exception was not thrown")
230+
} catch (e: Throwable) {
231+
return e
232+
}
233+
}
234+
235+
companion object {
236+
private val VALIDATION_EXCEPTION = ValidationException.builder()
237+
.message("Resource validation failed")
238+
.build()
239+
}
84240
}

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,6 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
9191
if (isProjectScope()) {
9292
LOG.debug {
9393
"Total size of source payload in KB: ${payloadContext.srcPayloadSize * 1.0 / TOTAL_BYTES_IN_KB} \n" +
94-
"Total size of build payload in KB: ${(payloadContext.buildPayloadSize ?: 0) * 1.0 / TOTAL_BYTES_IN_KB} \n" +
9594
"Total size of source zip file in KB: ${payloadContext.srcZipFileSize * 1.0 / TOTAL_BYTES_IN_KB} \n" +
9695
"Total number of lines reviewed: ${payloadContext.totalLines} \n" +
9796
"Total number of files included in payload: ${payloadContext.totalFiles} \n" +

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/sessionconfig/CodeScanSessionConfig.kt‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -334,7 +334,8 @@ data class PayloadContext(
334334
val scannedFiles: List<VirtualFile>,
335335
val srcPayloadSize: Long,
336336
val srcZipFileSize: Long,
337-
val buildPayloadSize: Long? = null,
337+
val payloadManifest: Set<Pair<String, Long>>? = null,
338+
val payloadLimitCrossed: Boolean? = false,
338339
)
339340

340341
data class PayloadMetadata(
@@ -343,4 +344,6 @@ data class PayloadMetadata(
343344
val linesScanned: Long,
344345
val language: CodewhispererLanguage,
345346
val codeDiff: String? = null,
347+
val payloadManifest: Set<Pair<String, Long>>? = null,
348+
val payloadLimitCrossed: Boolean? = false,
346349
)

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ open class CodeTestException(
1616
internal fun noFileOpenError(): Nothing =
1717
throw CodeTestException(message("codewhisperer.codescan.no_file_open"), "ProjectZipError")
1818

19-
internal fun fileTooLarge(): Nothing =
19+
fun fileTooLarge(): Nothing =
2020
throw CodeTestException(message("codewhisperer.codescan.file_too_large_telemetry"), "ProjectZipError")
2121

2222
internal fun cannotFindFile(errorMessage: String, filepath: String): Nothing =

0 commit comments

Comments
 (0)