Skip to content

Commit aa1a1c4

Browse files
tverneytverney
committed
fix(amazonqFeatureDev): include stop generation
1 parent b88f7d8 commit aa1a1c4

File tree

18 files changed

+1936
-1740
lines changed

18 files changed

+1936
-1740
lines changed

‎package-lock.json‎

Lines changed: 6 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClient.kt‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTo
3636
import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata
3737
import software.aws.toolkits.jetbrains.settings.AwsSettings
3838
import java.time.Instant
39+
import java.util.UUID
3940
import software.amazon.awssdk.services.codewhispererruntime.model.ChatTriggerType as SyncChatTriggerType
4041

4142
@Service(Service.Level.PROJECT)
@@ -87,9 +88,10 @@ class FeatureDevClient(private val project: Project) {
8788
CreateTaskAssistConversationRequest.builder().build()
8889
)
8990

90-
fun createTaskAssistUploadUrl(conversationId: String, contentChecksumSha256: String, contentLength: Long): CreateUploadUrlResponse =
91+
fun createTaskAssistUploadUrl(conversationId: String, contentChecksumSha256: String, contentLength: Long, uploadId: String): CreateUploadUrlResponse =
9192
bearerClient().createUploadUrl {
9293
it.contentChecksumType(ContentChecksumType.SHA_256)
94+
.uploadId(uploadId)
9395
.contentChecksum(contentChecksumSha256)
9496
.contentLength(contentLength)
9597
.artifactType(ArtifactType.SOURCE_CODE)
@@ -105,7 +107,7 @@ class FeatureDevClient(private val project: Project) {
105107
)
106108
}
107109

108-
fun startTaskAssistCodeGeneration(conversationId: String, uploadId: String, userMessage: String): StartTaskAssistCodeGenerationResponse = bearerClient()
110+
fun startTaskAssistCodeGeneration(conversationId: String, uploadId: String, userMessage: String, codeGenerationId: UUID, currentCodeGenerationId: UUID?): StartTaskAssistCodeGenerationResponse = bearerClient()
109111
.startTaskAssistCodeGeneration {
110112
request ->
111113
request
@@ -119,7 +121,7 @@ class FeatureDevClient(private val project: Project) {
119121
it
120122
.programmingLanguage { pl -> pl.languageName("javascript") } // This parameter is omitted by featureDev but required in the request
121123
.uploadId(uploadId)
122-
}
124+
}.codeGenerationId(codeGenerationId.toString()).currentCodeGenerationId(currentCodeGenerationId.toString())
123125
}
124126

125127
fun getTaskAssistCodeGeneration(conversationId: String, codeGenerationId: String): GetTaskAssistCodeGenerationResponse = bearerClient()

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -495,6 +495,7 @@ class FeatureDevController(
495495
}
496496
}
497497

498+
498499
private suspend fun handleChat(
499500
tabId: String,
500501
message: String,

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationState.kt‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.ge
1717
import software.aws.toolkits.resources.message
1818
import software.aws.toolkits.telemetry.AmazonqTelemetry
1919
import software.aws.toolkits.telemetry.Result
20+
import java.util.UUID
2021

2122
private val logger = getLogger<CodeGenerationState>()
2223

@@ -29,7 +30,8 @@ class CodeGenerationState(
2930
val repositorySize: Double,
3031
val messenger: MessagePublisher,
3132
var codeGenerationRemainingIterationCount: Int? = null,
32-
var codeGenerationTotalIterationCount: Int? = null
33+
var codeGenerationTotalIterationCount: Int? = null,
34+
var currentCodeGenerationId: UUID? = null
3335
) : SessionState {
3436
override val phase = SessionStatePhase.CODEGEN
3537

@@ -41,10 +43,13 @@ class CodeGenerationState(
4143
var numberOfReferencesGenerated: Int? = null
4244
var numberOfFilesGenerated: Int? = null
4345
try {
46+
val codeGenerationId = UUID.randomUUID()
4447
val response = config.featureDevService.startTaskAssistCodeGeneration(
4548
conversationId = config.conversationId,
4649
uploadId = uploadId,
47-
message = action.msg
50+
message = action.msg,
51+
codeGenerationId = codeGenerationId,
52+
currentCodeGenerationId = currentCodeGenerationId
4853
)
4954

5055
messenger.sendAnswerPart(

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationState.kt‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import software.aws.toolkits.resources.message
1616
import software.aws.toolkits.telemetry.AmazonqTelemetry
1717
import software.aws.toolkits.telemetry.AmazonqUploadIntent
1818
import software.aws.toolkits.telemetry.Result
19+
import java.util.UUID
1920

2021
private val logger = getLogger<PrepareCodeGenerationState>()
2122

@@ -49,16 +50,18 @@ class PrepareCodeGenerationState(
4950
zipFileLength = repoZipResult.contentLength
5051
val fileToUpload = repoZipResult.payload
5152

53+
val uploadId = UUID.randomUUID();
5254
val uploadUrlResponse = config.featureDevService.createUploadUrl(
5355
config.conversationId,
5456
zipFileChecksum,
55-
zipFileLength
57+
zipFileLength,
58+
uploadId
5659
)
5760

5861
uploadArtifactToS3(uploadUrlResponse.uploadUrl(), fileToUpload, zipFileChecksum, zipFileLength, uploadUrlResponse.kmsKeyArn())
5962
deleteUploadArtifact(fileToUpload)
6063

61-
this.uploadId = uploadUrlResponse.uploadId()
64+
this.uploadId = uploadId.toString()
6265
messenger.sendAnswerPart(tabId = this.tabID, message = message("amazonqFeatureDev.placeholder.context_gathering_complete"))
6366
messenger.sendUpdatePlaceholder(tabId = this.tabID, newPlaceholder = message("amazonqFeatureDev.placeholder.context_gathering_complete"))
6467
nextState = CodeGenerationState(

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevService.kt‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.ge
3131
import software.aws.toolkits.resources.message
3232
import software.aws.toolkits.telemetry.AmazonqTelemetry
3333
import software.aws.toolkits.telemetry.Result
34+
import java.util.UUID
3435

3536
private val logger = getLogger<FeatureDevClient>()
3637

@@ -76,17 +77,18 @@ class FeatureDevService(val proxyClient: FeatureDevClient, val project: Project)
7677
}
7778
}
7879

79-
fun createUploadUrl(conversationId: String, contentChecksumSha256: String, contentLength: Long):
80+
fun createUploadUrl(conversationId: String, contentChecksumSha256: String, contentLength: Long, uploadId: UUID):
8081
CreateUploadUrlResponse {
8182
try {
8283
logger.debug { "Executing createUploadUrl with conversationId $conversationId" }
8384
val uploadUrlResponse = proxyClient.createTaskAssistUploadUrl(
8485
conversationId,
8586
contentChecksumSha256,
86-
contentLength
87+
contentLength,
88+
uploadId.toString()
8789
)
8890
logger.debug {
89-
"$FEATURE_NAME: Created upload url: {uploadId: ${uploadUrlResponse.uploadId()}, requestId: ${uploadUrlResponse.responseMetadata().requestId()}}"
91+
"$FEATURE_NAME: Created upload url: {uploadId: ${uploadId.toString()}, requestId: ${uploadUrlResponse.responseMetadata().requestId()}}"
9092
}
9193
return uploadUrlResponse
9294
} catch (e: Exception) {
@@ -105,14 +107,16 @@ class FeatureDevService(val proxyClient: FeatureDevClient, val project: Project)
105107
}
106108
}
107109

108-
fun startTaskAssistCodeGeneration(conversationId: String, uploadId: String, message: String):
110+
fun startTaskAssistCodeGeneration(conversationId: String, uploadId: String, message: String, codeGenerationId: UUID, currentCodeGenerationId: UUID?):
109111
StartTaskAssistCodeGenerationResponse {
110112
try {
111113
logger.debug { "Executing startTaskAssistCodeGeneration with conversationId: $conversationId , uploadId: $uploadId" }
112114
val startCodeGenerationResponse = proxyClient.startTaskAssistCodeGeneration(
113115
conversationId,
114116
uploadId,
115-
message
117+
message,
118+
codeGenerationId,
119+
currentCodeGenerationId
116120
)
117121

118122
logger.debug { "$FEATURE_NAME: Started code generation with requestId: ${startCodeGenerationResponse.responseMetadata().requestId()}" }

‎plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevTestBase.kt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRul
4242
import software.aws.toolkits.jetbrains.utils.rules.addModule
4343
import java.io.File
4444
import java.time.Instant
45+
import java.util.UUID
4546

4647
open class FeatureDevTestBase(
4748
@Rule @JvmField
@@ -59,6 +60,8 @@ open class FeatureDevTestBase(
5960
internal val testRequestId = "test_aws_request_id"
6061
internal val testConversationId = "1234"
6162
internal val userMessage = "test-user-message"
63+
internal val codeGenerationId = UUID.randomUUID()
64+
internal val currentCodeGenerationId = UUID.randomUUID()
6265
internal val testUploadId = "5678"
6366
internal val testRepositorySize = 20.0 // Picked a random size
6467
internal val testCodeGenerationId = "1234"
@@ -69,6 +72,7 @@ open class FeatureDevTestBase(
6972
internal val testReferences = listOf(CodeReferenceGenerated())
7073
internal val testChecksumSha = "test-sha"
7174
internal val testContentLength: Long = 40
75+
internal val uploadId = UUID.randomUUID()
7276

7377
internal val exampleSendTelemetryEventResponse = SendTelemetryEventResponse.builder()
7478
.responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))

‎plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationStateTest.kt‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ class CodeGenerationStateTest : FeatureDevTestBase() {
7171
fun `test code generated is complete`() {
7272
val action = SessionStateAction("test-task", userMessage)
7373
every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleCompleteGetTaskAssistCodeGenerationResponse
74-
every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
74+
every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
7575
coEvery { featureDevService.exportTaskAssistArchiveResult(any()) } returns
7676
CodeGenerationStreamResult(testFilePaths, testDeletedFiles, testReferences)
7777

@@ -95,26 +95,26 @@ class CodeGenerationStateTest : FeatureDevTestBase() {
9595
}
9696
assertThat(codeGenerationState.phase).isEqualTo(SessionStatePhase.CODEGEN)
9797
coVerify(exactly = 1) { messenger.sendAnswerPart(testTabId, message("amazonqFeatureDev.code_generation.generating_code")) }
98-
verify(exactly = 1) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage) }
98+
verify(exactly = 1) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage, codeGenerationId, currentCodeGenerationId) }
9999
verify(exactly = 1) { featureDevService.getTaskAssistCodeGeneration(testConversationId, testCodeGenerationId) }
100100
coVerify(exactly = 1) { featureDevService.exportTaskAssistArchiveResult(testConversationId) }
101101
}
102102

103103
@Test(expected = FeatureDevException::class)
104104
fun `test code generation failed`() = runTest {
105-
every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
105+
every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
106106
every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleFailedGetTaskAssistCodeGenerationResponse
107107

108108
codeGenerationState.interact(action)
109109

110-
verify(exactly = 1) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage) }
110+
verify(exactly = 1) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage, codeGenerationId, currentCodeGenerationId) }
111111
verify(exactly = 1) { featureDevService.getTaskAssistCodeGeneration(testConversationId, testCodeGenerationId) }
112112
coVerify(exactly = 0) { featureDevService.exportTaskAssistArchiveResult(any()) }
113113
}
114114

115115
@Test
116116
fun `test code generation returns any other handled status`() {
117-
every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
117+
every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
118118
every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleOtherGetTaskAssistCodeGenerationResponse
119119

120120
assertThatThrownBy {
@@ -123,14 +123,14 @@ class CodeGenerationStateTest : FeatureDevTestBase() {
123123
}
124124
}.isExactlyInstanceOf(IllegalStateException::class.java).withFailMessage("Unknown status: $otherStatus")
125125

126-
verify(exactly = 1) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage) }
126+
verify(exactly = 1) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage, codeGenerationId, currentCodeGenerationId) }
127127
verify(exactly = 1) { featureDevService.getTaskAssistCodeGeneration(testConversationId, testCodeGenerationId) }
128128
coVerify(exactly = 0) { featureDevService.exportTaskAssistArchiveResult(any()) }
129129
}
130130

131131
@Test
132132
fun `test code generation returns in progress at least once`() {
133-
every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
133+
every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
134134
every {
135135
featureDevService.getTaskAssistCodeGeneration(any(), any())
136136
} returnsMany listOf(exampleGetTaskAssistConversationResponse, exampleCompleteGetTaskAssistCodeGenerationResponse)
@@ -140,14 +140,14 @@ class CodeGenerationStateTest : FeatureDevTestBase() {
140140
codeGenerationState.interact(action)
141141
}
142142

143-
verify(exactly = 1) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage) }
143+
verify(exactly = 1) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage, codeGenerationId, currentCodeGenerationId) }
144144
verify(exactly = 2) { featureDevService.getTaskAssistCodeGeneration(testConversationId, testCodeGenerationId) }
145145
coVerify(exactly = 1) { featureDevService.exportTaskAssistArchiveResult(testConversationId) }
146146
}
147147

148148
@Test
149149
fun `test using all polling count`() {
150-
every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
150+
every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
151151
every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleGetTaskAssistConversationResponse
152152

153153
runTest {
@@ -163,7 +163,7 @@ class CodeGenerationStateTest : FeatureDevTestBase() {
163163
assertThat(actual.interaction.content).isEqualTo("")
164164
}
165165

166-
verify(exactly = 1) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage) }
166+
verify(exactly = 1) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage, codeGenerationId, currentCodeGenerationId) }
167167
verify(exactly = 180) { featureDevService.getTaskAssistCodeGeneration(testConversationId, testCodeGenerationId) }
168168
coVerify(exactly = 0) { featureDevService.exportTaskAssistArchiveResult(testConversationId) }
169169
}

‎plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationStateTest.kt‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ class PrepareCodeGenerationStateTest : FeatureDevTestBase() {
6666
every { deleteUploadArtifact(any()) } just runs
6767

6868
every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleCompleteGetTaskAssistCodeGenerationResponse
69-
every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
69+
every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
7070
coEvery { featureDevService.exportTaskAssistArchiveResult(any()) } returns exampleExportTaskAssistResultArchiveResponse
7171

7272
mockkStatic(MessagePublisher::sendAnswerPart)
@@ -85,7 +85,7 @@ class PrepareCodeGenerationStateTest : FeatureDevTestBase() {
8585
val action = SessionStateAction("test-task", userMessage)
8686

8787
whenever(repoContext.getProjectZip()).thenReturn(repoZipResult)
88-
every { featureDevService.createUploadUrl(any(), any(), any()) } returns exampleCreateUploadUrlResponse
88+
every { featureDevService.createUploadUrl(any(), any(), any(), any()) } returns exampleCreateUploadUrlResponse
8989

9090
runTest {
9191
val actual = prepareCodeGenerationState.interact(action)
@@ -94,6 +94,6 @@ class PrepareCodeGenerationStateTest : FeatureDevTestBase() {
9494
}
9595
assertThat(prepareCodeGenerationState.phase).isEqualTo(SessionStatePhase.CODEGEN)
9696
verify(repoContext, times(1)).getProjectZip()
97-
io.mockk.verify(exactly = 1) { featureDevService.createUploadUrl(testConversationId, testChecksumSha, testContentLength) }
97+
io.mockk.verify(exactly = 1) { featureDevService.createUploadUrl(testConversationId, testChecksumSha, testContentLength, uploadId) }
9898
}
9999
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ object CodeWhispererConstants {
105105
}
106106

107107
object Config {
108-
const val CODEWHISPERER_ENDPOINT = "https://codewhisperer.us-east-1.amazonaws.com/" // PROD
108+
const val CODEWHISPERER_ENDPOINT = "https://rts.gamma-us-east-1.codewhisperer.ai.aws.dev/" // PROD
109109
const val CODEWHISPERER_IDPOOL_ID = "us-east-1:70717e99-906f-4add-908c-bd9074a2f5b9"
110110
val Sigv4ClientRegion = Region.US_EAST_1
111111
val BearerClientRegion = Region.US_EAST_1

0 commit comments

Comments
 (0)