Skip to content

Commit 556f3f1

Browse files
ashishrp-awsashishrp-aws
authored
Merge branch 'main' into workspace-test
2 parents 346cee3 + 52d7701 commit 556f3f1

File tree

25 files changed

+384
-123
lines changed

25 files changed

+384
-123
lines changed

‎.changes/next-release/bugfix-2df57233-310e-4788-be6c-f3b37767e60a.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "bugfix",
3+
"description" : "Amazon Q can update mvn and gradle build files"
4+
}

‎.changes/next-release/bugfix-693e3f08-15f7-4453-b304-e473c58bf46b.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "bugfix",
3+
"description" : "Amazon Q /dev: Remove hard-coded limits and instead rely server-side data to communicate number of code generations remaining"
4+
}

‎.changes/next-release/feature-ccbbe2b3-2bad-443b-9cf1-7708eb6f77d9.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "feature",
3+
"description" : "Enhance Q inline completion context fetching for better suggestion quality"
4+
}

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

Lines changed: 30 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Delete
2828
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.NewFileZipInfo
2929
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.PrepareCodeGenerationState
3030
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Session
31-
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionState
3231
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.CancellationTokenSource
3332
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.InsertAction
3433
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions
@@ -64,7 +63,7 @@ suspend fun FeatureDevController.onCodeGeneration(
6463
var totalIterations: Int? = state.codeGenerationTotalIterationCount
6564

6665
if (state.token?.token?.isCancellationRequested() == true) {
67-
disposeToken(state, messenger, tabId, state.currentIteration?.let { CODE_GENERATION_RETRY_LIMIT.minus(it) }, CODE_GENERATION_RETRY_LIMIT)
66+
disposeToken(messenger, tabId, state.codeGenerationRemainingIterationCount, state.codeGenerationTotalIterationCount)
6867
return
6968
}
7069

@@ -96,7 +95,7 @@ suspend fun FeatureDevController.onCodeGeneration(
9695
}
9796

9897
if (state.token?.token?.isCancellationRequested() == true) {
99-
disposeToken(state, messenger, tabId, state.currentIteration?.let { CODE_GENERATION_RETRY_LIMIT.minus(it) }, CODE_GENERATION_RETRY_LIMIT)
98+
disposeToken(messenger, tabId, state.codeGenerationRemainingIterationCount, state.codeGenerationTotalIterationCount)
10099
return
101100
}
102101

@@ -133,14 +132,20 @@ suspend fun FeatureDevController.onCodeGeneration(
133132
tabId = tabId,
134133
messageType = FeatureDevMessageType.Answer,
135134
message =
136-
if (remainingIterations == 0) {
137-
message("amazonqFeatureDev.code_generation.iteration_zero")
138-
} else {
135+
if (remainingIterations > 2) {
136+
message("amazonqFeatureDev.code_generation.iteration_counts_ask_to_add_code_or_feedback")
137+
} else if (remainingIterations > 0) {
139138
message(
140139
"amazonqFeatureDev.code_generation.iteration_counts",
141140
remainingIterations,
142141
totalIterations,
143142
)
143+
} else {
144+
message(
145+
"amazonqFeatureDev.code_generation.iteration_counts_ask_to_add_code",
146+
remainingIterations,
147+
totalIterations,
148+
)
144149
},
145150
)
146151
}
@@ -196,7 +201,6 @@ suspend fun FeatureDevController.onCodeGeneration(
196201
}
197202

198203
private suspend fun disposeToken(
199-
state: SessionState,
200204
messenger: MessagePublisher,
201205
tabId: String,
202206
remainingIterations: Int?,
@@ -234,16 +238,25 @@ private suspend fun disposeToken(
234238
return
235239
}
236240

237-
messenger.sendAnswer(
238-
tabId = tabId,
239-
messageType = FeatureDevMessageType.Answer,
240-
message =
241-
message(
242-
"amazonqFeatureDev.code_generation.stopped_code_generation",
243-
remainingIterations ?: state.currentIteration?.let { CODE_GENERATION_RETRY_LIMIT - it } as Any,
244-
totalIterations ?: CODE_GENERATION_RETRY_LIMIT,
245-
),
246-
)
241+
if (remainingIterations !== null && totalIterations !== null && remainingIterations <= 2) {
242+
messenger.sendAnswer(
243+
tabId = tabId,
244+
messageType = FeatureDevMessageType.Answer,
245+
message =
246+
message(
247+
"amazonqFeatureDev.code_generation.stopped_code_generation",
248+
remainingIterations,
249+
totalIterations,
250+
),
251+
)
252+
} else {
253+
messenger.sendAnswer(
254+
tabId = tabId,
255+
messageType = FeatureDevMessageType.Answer,
256+
message =
257+
message("amazonqFeatureDev.code_generation.stopped_code_generation_no_iteration_count_display"),
258+
)
259+
}
247260

248261
messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = true)
249262

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

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,12 @@ class CodeGenerationState(
8484
numberOfFilesGenerated = codeGenerationResult.newFiles.size
8585
codeGenerationRemainingIterationCount = codeGenerationResult.codeGenerationRemainingIterationCount
8686
codeGenerationTotalIterationCount = codeGenerationResult.codeGenerationTotalIterationCount
87+
currentIteration =
88+
if (codeGenerationRemainingIterationCount != null && codeGenerationTotalIterationCount != null) {
89+
codeGenerationTotalIterationCount?.let { total -> codeGenerationRemainingIterationCount?.let { remaining -> total - remaining } }
90+
} else {
91+
currentIteration?.plus(1)
92+
}
8793

8894
runCatching {
8995
var insertedLines = 0
@@ -128,7 +134,7 @@ class CodeGenerationState(
128134
filePaths = codeGenerationResult.newFiles,
129135
deletedFiles = codeGenerationResult.deletedFiles,
130136
references = codeGenerationResult.references,
131-
currentIteration = currentIteration?.plus(1),
137+
currentIteration = currentIteration,
132138
uploadId = uploadId,
133139
messenger = messenger,
134140
codeGenerationRemainingIterationCount = codeGenerationRemainingIterationCount,
@@ -182,17 +188,22 @@ private suspend fun CodeGenerationState.generateCode(
182188
): CodeGenerationResult {
183189
val pollCount = 360
184190
val requestDelay = 5000L
191+
var codeGenerationRemainingIterationCount: Int? = null
192+
var codeGenerationTotalIterationCount: Int? = null
185193

186194
repeat(pollCount) {
187195
if (token?.token?.isCancellationRequested() == true) {
188-
return CodeGenerationResult(emptyList(), emptyList(), emptyList())
196+
return CodeGenerationResult(emptyList(), emptyList(), emptyList(), codeGenerationRemainingIterationCount, codeGenerationTotalIterationCount)
189197
}
190198
val codeGenerationResultState =
191199
config.featureDevService.getTaskAssistCodeGeneration(
192200
conversationId = config.conversationId,
193201
codeGenerationId = codeGenerationId,
194202
)
195203

204+
codeGenerationRemainingIterationCount = codeGenerationResultState.codeGenerationRemainingIterationCount()
205+
codeGenerationTotalIterationCount = codeGenerationResultState.codeGenerationTotalIterationCount()
206+
196207
when (codeGenerationResultState.codeGenerationStatus().status()) {
197208
CodeGenerationWorkflowStatus.COMPLETE -> {
198209
val codeGenerationStreamResult =
@@ -207,8 +218,8 @@ private suspend fun CodeGenerationState.generateCode(
207218
newFiles = newFileInfo,
208219
deletedFiles = deletedFileInfo,
209220
references = codeGenerationStreamResult.references,
210-
codeGenerationRemainingIterationCount = codeGenerationResultState.codeGenerationRemainingIterationCount(),
211-
codeGenerationTotalIterationCount = codeGenerationResultState.codeGenerationTotalIterationCount(),
221+
codeGenerationRemainingIterationCount = codeGenerationRemainingIterationCount,
222+
codeGenerationTotalIterationCount = codeGenerationTotalIterationCount,
212223
)
213224
}
214225
CodeGenerationWorkflowStatus.IN_PROGRESS -> {
@@ -252,7 +263,7 @@ private suspend fun CodeGenerationState.generateCode(
252263
}
253264
}
254265

255-
return CodeGenerationResult(emptyList(), emptyList(), emptyList())
266+
return CodeGenerationResult(emptyList(), emptyList(), emptyList(), codeGenerationRemainingIterationCount, codeGenerationTotalIterationCount)
256267
}
257268

258269
fun registerNewFiles(newFileContents: Map<String, String>): List<NewFileZipInfo> =

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@ class ConversationNotStartedState(
99
override var approach: String,
1010
override val tabID: String,
1111
override var token: CancellationTokenSource?,
12-
override var codeGenerationRemainingIterationCount: Int?,
13-
override var codeGenerationTotalIterationCount: Int?,
12+
override var codeGenerationRemainingIterationCount: Int? = null,
13+
override var codeGenerationTotalIterationCount: Int? = null,
1414
override var currentIteration: Int?,
1515
override var diffMetricsProcessed: DiffMetricsProcessed,
1616
) : SessionState {

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,8 +55,6 @@ class Session(val tabID: String, val project: Project) {
5555
approach = "",
5656
tabID = tabID,
5757
token = null,
58-
codeGenerationRemainingIterationCount = 0,
59-
codeGenerationTotalIterationCount = CODE_GENERATION_RETRY_LIMIT,
6058
currentIteration = 0,
6159
diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet())
6260
)

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/editor/context/focusArea/FocusAreaContextExtractor.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,7 @@ class FocusAreaContextExtractor(private val fqnWebviewAdapter: FqnWebviewAdapter
180180
val startOffset = 0.coerceAtLeast(offset - halfMaxCharacters)
181181
val endOffset = fileText.length.coerceAtMost(offset + halfMaxCharacters)
182182

183-
// Adjust the start and end offsets if necessary to ensure a total of 10k characters
183+
// Adjust the start and end offsets if necessary to ensure a total of 40k characters
184184
val excessCharacters = maxCharacters - (endOffset - startOffset)
185185
val adjustedStartOffset = 0.coerceAtLeast(startOffset - excessCharacters)
186186
val adjustedEndOffset = fileText.length.coerceAtMost(endOffset + excessCharacters)
@@ -198,7 +198,7 @@ class FocusAreaContextExtractor(private val fqnWebviewAdapter: FqnWebviewAdapter
198198
}
199199

200200
companion object {
201-
const val MAX_LENGTH = 10000
201+
const val MAX_LENGTH = 40000
202202
}
203203
}
204204

‎plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/workspace/context/ProjectContextProviderTest.kt‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.project.EncoderServer
4141
import software.aws.toolkits.jetbrains.services.amazonq.project.IndexRequest
4242
import software.aws.toolkits.jetbrains.services.amazonq.project.IndexUpdateMode
4343
import software.aws.toolkits.jetbrains.services.amazonq.project.InlineBm25Chunk
44+
import software.aws.toolkits.jetbrains.services.amazonq.project.InlineContextTarget
4445
import software.aws.toolkits.jetbrains.services.amazonq.project.LspMessage
4546
import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextProvider
4647
import software.aws.toolkits.jetbrains.services.amazonq.project.QueryChatRequest
@@ -237,13 +238,13 @@ class ProjectContextProviderTest {
237238
@Test
238239
fun `queryInline should send correct encrypted request to lsp`() = runTest {
239240
sut = ProjectContextProvider(project, encoderServer, this)
240-
sut.queryInline("foo", "Foo.java")
241+
sut.queryInline("foo", "Foo.java", InlineContextTarget.CODEMAP)
241242
advanceUntilIdle()
242243

243-
val request = QueryInlineCompletionRequest("foo", "Foo.java")
244+
val request = QueryInlineCompletionRequest("foo", "Foo.java", "codemap")
244245
val requestJson = mapper.writeValueAsString(request)
245246

246-
assertThat(mapper.readTree(requestJson)).isEqualTo(mapper.readTree("""{ "query": "foo", "filePath": "Foo.java" }"""))
247+
assertThat(mapper.readTree(requestJson)).isEqualTo(mapper.readTree("""{ "query": "foo", "filePath": "Foo.java", "target": "codemap" }"""))
247248

248249
val encryptedRequest = encoderServer.encrypt(requestJson)
249250
wireMock.verify(
@@ -315,7 +316,7 @@ class ProjectContextProviderTest {
315316
)
316317

317318
assertThrows<Exception> {
318-
sut.queryInline("foo", "filepath")
319+
sut.queryInline("foo", "filepath", InlineContextTarget.CODEMAP)
319320
advanceUntilIdle()
320321
}
321322
}
@@ -326,7 +327,7 @@ class ProjectContextProviderTest {
326327
fun `query inline should return deserialized bm25 chunks`() = runTest {
327328
sut = ProjectContextProvider(project, encoderServer, this)
328329
advanceUntilIdle()
329-
val r = sut.queryInline("foo", "filepath")
330+
val r = sut.queryInline("foo", "filepath", InlineContextTarget.CODEMAP)
330331
assertThat(r).hasSize(3)
331332
assertThat(r[0]).isEqualTo(
332333
InlineBm25Chunk(
@@ -374,7 +375,7 @@ class ProjectContextProviderTest {
374375

375376
// it won't throw if it's executed within TestDispatcher context
376377
withContext(getCoroutineBgContext()) {
377-
sut.queryInline("foo", "bar")
378+
sut.queryInline("foo", "bar", InlineContextTarget.CODEMAP)
378379
}
379380

380381
advanceUntilIdle()

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

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,17 @@ import org.mockito.kotlin.whenever
1313
import software.aws.toolkits.jetbrains.services.amazonq.FeatureDevSessionContext
1414
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FeatureDevTestBase
1515
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.FeatureDevService
16+
import software.aws.toolkits.jetbrains.utils.rules.HeavyJavaCodeInsightTestFixtureRule
17+
import software.aws.toolkits.jetbrains.utils.rules.addFileToModule
18+
import java.util.zip.ZipFile
19+
20+
class FeatureDevSessionContextTest : FeatureDevTestBase(HeavyJavaCodeInsightTestFixtureRule()) {
21+
22+
private fun addFilesToProjectModule(vararg path: String) {
23+
val module = projectRule.module
24+
path.forEach { projectRule.fixture.addFileToModule(module, it, it) }
25+
}
1626

17-
class FeatureDevSessionContextTest : FeatureDevTestBase() {
1827
@Rule
1928
@JvmField
2029
val ruleChain = RuleChain(projectRule, disposableRule)
@@ -40,13 +49,64 @@ class FeatureDevSessionContextTest : FeatureDevTestBase() {
4049
fun testWithValidFile() {
4150
val ktFile = mock<VirtualFile>()
4251
whenever(ktFile.extension).thenReturn("kt")
52+
whenever(ktFile.path).thenReturn("code.kt")
4353
assertTrue(featureDevSessionContext.isFileExtensionAllowed(ktFile))
4454
}
4555

4656
@Test
4757
fun testWithInvalidFile() {
4858
val txtFile = mock<VirtualFile>()
4959
whenever(txtFile.extension).thenReturn("txt")
60+
whenever(txtFile.path).thenReturn("file.txt")
5061
assertFalse(featureDevSessionContext.isFileExtensionAllowed(txtFile))
5162
}
63+
64+
@Test
65+
fun testAllowedFilePath() {
66+
val allowedPaths = listOf("build.gradle", "gradle.properties", ".mvn/wrapper/maven-wrapper.properties")
67+
allowedPaths.forEach({
68+
val txtFile = mock<VirtualFile>()
69+
whenever(txtFile.path).thenReturn(it)
70+
whenever(txtFile.extension).thenReturn(it.split(".").last())
71+
assertTrue(featureDevSessionContext.isFileExtensionAllowed(txtFile))
72+
})
73+
}
74+
75+
@Test
76+
fun testZipProject() {
77+
addFilesToProjectModule(
78+
".gradle/cached.jar",
79+
"src/MyClass.java",
80+
"gradlew",
81+
"gradlew.bat",
82+
"README.md",
83+
"settings.gradle",
84+
"build.gradle",
85+
"gradle/wrapper/gradle-wrapper.properties",
86+
)
87+
88+
val zipResult = featureDevSessionContext.getProjectZip()
89+
val zipPath = zipResult.payload.path
90+
91+
val zippedFiles = mutableSetOf<String>()
92+
ZipFile(zipPath).use { zipFile ->
93+
for (entry in zipFile.entries()) {
94+
if (!entry.name.endsWith("/")) {
95+
zippedFiles.add(entry.name)
96+
}
97+
}
98+
}
99+
100+
val expectedFiles = setOf(
101+
"src/MyClass.java",
102+
"gradlew",
103+
"gradlew.bat",
104+
"README.md",
105+
"settings.gradle",
106+
"build.gradle",
107+
"gradle/wrapper/gradle-wrapper.properties",
108+
)
109+
110+
assertTrue(zippedFiles == expectedFiles)
111+
}
52112
}

0 commit comments

Comments
 (0)