Skip to content

Commit 6a83cfc

Browse files
aws-toolkit-automationaws-toolkit-automation
authored
Merge main into feature/gettingstarted
2 parents a3e2a0a + cb08c9a commit 6a83cfc

File tree

3 files changed

+97
-118
lines changed

3 files changed

+97
-118
lines changed

‎.changes/next-release/bugfix-05445066-8f16-47fc-be95-503f9a5f0743.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" : "CodeWhisperer: Fix an issue where an IndexOutOfBoundException could be thrown when using CodeWhisperer"
4+
}

‎jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererRecommendationManager.kt‎

Lines changed: 36 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -3,105 +3,48 @@
33

44
package software.aws.toolkits.jetbrains.services.codewhisperer.service
55

6-
import com.intellij.openapi.command.WriteCommandAction
76
import com.intellij.openapi.components.service
8-
import com.intellij.openapi.editor.RangeMarker
9-
import com.intellij.openapi.util.TextRange
10-
import com.intellij.psi.PsiDocumentManager
11-
import com.intellij.psi.PsiFileFactory
12-
import com.intellij.psi.codeStyle.CodeStyleManager
13-
import com.intellij.util.LocalTimeCounter
147
import software.amazon.awssdk.services.codewhispererruntime.model.Completion
15-
import software.amazon.awssdk.services.codewhispererruntime.model.Reference
168
import software.amazon.awssdk.services.codewhispererruntime.model.Span
179
import software.aws.toolkits.jetbrains.services.codewhisperer.model.DetailContext
1810
import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationChunk
1911
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getCompletionType
2012
import kotlin.math.max
13+
import kotlin.math.min
2114

2215
class CodeWhispererRecommendationManager {
23-
fun reformat(requestContext: RequestContext, recommendation: Completion): Completion {
24-
val project = requestContext.project
25-
val editor = requestContext.editor
26-
val document = editor.document
27-
16+
fun reformatReference(requestContext: RequestContext, recommendation: Completion): Completion {
2817
// startOffset is the offset at the start of user input since invocation
2918
val invocationStartOffset = requestContext.caretPosition.offset
30-
val startOffsetSinceUserInput = editor.caretModel.offset
31-
32-
// Create a temp file for capturing reformatted text and updated content spans
33-
val tempPsiFile = PsiDocumentManager.getInstance(project).getPsiFile(document)?.let { psiFile ->
34-
PsiFileFactory.getInstance(project).createFileFromText(
35-
"codewhisperer_temp",
36-
psiFile.fileType,
37-
document.text,
38-
LocalTimeCounter.currentTime(),
39-
true
40-
)
41-
}
42-
val tempDocument = tempPsiFile?.let { psiFile ->
43-
PsiDocumentManager.getInstance(project).getDocument(psiFile)
44-
} ?: return recommendation
4519

20+
val startOffsetSinceUserInput = requestContext.editor.caretModel.offset
4621
val endOffset = invocationStartOffset + recommendation.content().length
22+
4723
if (startOffsetSinceUserInput > endOffset) return recommendation
48-
WriteCommandAction.runWriteCommandAction(project) {
49-
tempDocument.insertString(invocationStartOffset, recommendation.content())
50-
PsiDocumentManager.getInstance(project).commitDocument(tempDocument)
51-
}
52-
val rangeMarkers = mutableMapOf<RangeMarker, Reference>()
53-
recommendation.references().forEach {
54-
val referenceStart = invocationStartOffset + it.recommendationContentSpan().start()
55-
if (referenceStart >= endOffset) return@forEach
56-
val tempEnd = invocationStartOffset + it.recommendationContentSpan().end()
57-
val referenceEnd = if (tempEnd <= endOffset) tempEnd else endOffset
58-
rangeMarkers[
59-
tempDocument.createRangeMarker(
60-
referenceStart,
61-
referenceEnd
62-
)
63-
] = it
64-
}
65-
val tempRangeMarker = tempDocument.createRangeMarker(invocationStartOffset, endOffset)
6624

67-
// Currently, only reformat(adjust line indent) starting from user's input
68-
WriteCommandAction.runWriteCommandAction(project) {
69-
CodeStyleManager.getInstance(project).adjustLineIndent(tempPsiFile, TextRange(startOffsetSinceUserInput, endOffset))
25+
val reformattedReferences = recommendation.references().filter {
26+
val referenceStart = invocationStartOffset + it.recommendationContentSpan().start()
27+
val referenceEnd = invocationStartOffset + it.recommendationContentSpan().end()
28+
referenceStart < endOffset && referenceEnd > startOffsetSinceUserInput
29+
}.map {
30+
val referenceStart = invocationStartOffset + it.recommendationContentSpan().start()
31+
val referenceEnd = invocationStartOffset + it.recommendationContentSpan().end()
32+
val updatedReferenceStart = max(referenceStart, startOffsetSinceUserInput)
33+
val updatedReferenceEnd = min(referenceEnd, endOffset)
34+
it.toBuilder().recommendationContentSpan(
35+
Span.builder()
36+
.start(updatedReferenceStart - invocationStartOffset)
37+
.end(updatedReferenceEnd - invocationStartOffset)
38+
.build()
39+
).build()
7040
}
7141

72-
val reformattedRecommendation = tempDocument.getText(TextRange(tempRangeMarker.startOffset, tempRangeMarker.endOffset))
73-
74-
val reformattedReferences = rangeMarkers.map { (rangeMarker, reference) ->
75-
reformatReference(reference, rangeMarker, invocationStartOffset)
76-
}
7742
return Completion.builder()
78-
.content(reformattedRecommendation)
43+
.content(recommendation.content())
7944
.references(reformattedReferences)
8045
.build()
8146
}
8247

83-
/**
84-
* Build new reference with updated contentSpan(start and end). Since it's reformatted, take the new start and
85-
* end from the rangeMarker which automatically tracks the range after reformatting
86-
*/
87-
fun reformatReference(originalReference: Reference, rangeMarker: RangeMarker, invocationStartOffset: Int): Reference {
88-
rangeMarker.apply {
89-
val documentContent = document.charsSequence
90-
91-
// has to plus 1 because right boundary is exclusive
92-
val spanEndOffset = documentContent.subSequence(0, endOffset).indexOfLast { char -> char != '\n' } + 1
93-
return originalReference
94-
.toBuilder()
95-
.recommendationContentSpan(
96-
Span.builder()
97-
.start(startOffset - invocationStartOffset)
98-
.end(spanEndOffset - invocationStartOffset)
99-
.build()
100-
)
101-
.build()
102-
}
103-
}
104-
10548
fun buildRecommendationChunks(
10649
recommendation: String,
10750
matchingSymbols: List<Pair<Int, Int>>
@@ -157,21 +100,33 @@ class CodeWhispererRecommendationManager {
157100
)
158101
}
159102

160-
val reformatted = reformat(requestContext, truncated)
161-
val isDiscardedByRightContextTruncationDedupe = !seen.add(reformatted.content())
103+
val isDiscardedByRightContextTruncationDedupe = !seen.add(truncated.content())
104+
val isDiscardedByBlankAfterTruncation = truncated.content().isBlank()
105+
if (isDiscardedByRightContextTruncationDedupe || isDiscardedByBlankAfterTruncation) {
106+
return@map DetailContext(
107+
requestId,
108+
it,
109+
truncated,
110+
isDiscarded = true,
111+
truncated.content().length != it.content().length,
112+
overlap,
113+
getCompletionType(it)
114+
)
115+
}
116+
val reformatted = reformatReference(requestContext, truncated)
162117
DetailContext(
163118
requestId,
164119
it,
165120
reformatted,
166-
isDiscardedByRightContextTruncationDedupe,
121+
isDiscarded = false,
167122
truncated.content().length != it.content().length,
168123
overlap,
169124
getCompletionType(it)
170125
)
171126
}
172127
}
173128

174-
private fun findRightContextOverlap(
129+
fun findRightContextOverlap(
175130
requestContext: RequestContext,
176131
recommendation: Completion
177132
): String {

‎jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererRecommendationManagerTest.kt‎

Lines changed: 57 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,21 @@
33

44
package software.aws.toolkits.jetbrains.services.codewhisperer
55

6-
import com.intellij.openapi.command.WriteCommandAction
6+
import com.intellij.openapi.application.ApplicationManager
77
import com.intellij.openapi.project.Project
88
import com.intellij.testFramework.DisposableRule
99
import com.intellij.testFramework.fixtures.CodeInsightTestFixture
10-
import com.intellij.testFramework.runInEdtAndGet
10+
import com.intellij.testFramework.replaceService
1111
import com.intellij.testFramework.runInEdtAndWait
1212
import org.assertj.core.api.Assertions.assertThat
1313
import org.junit.Before
1414
import org.junit.Rule
1515
import org.junit.Test
16-
import software.amazon.awssdk.services.codewhispererruntime.model.Reference
16+
import org.mockito.kotlin.any
17+
import org.mockito.kotlin.doReturn
18+
import org.mockito.kotlin.spy
19+
import org.mockito.kotlin.stub
20+
import software.aws.toolkits.core.utils.test.aString
1721
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererRecommendationManager
1822
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
1923

@@ -27,12 +31,9 @@ class CodeWhispererRecommendationManagerTest {
2731
val disposableRule = DisposableRule()
2832

2933
private val documentContentContent = "012345678"
34+
private lateinit var sut: CodeWhispererRecommendationManager
3035
private lateinit var fixture: CodeInsightTestFixture
3136
private lateinit var project: Project
32-
private val originalReference = Reference.builder()
33-
.licenseName("test_license")
34-
.repository("test_repo")
35-
.build()
3637

3738
@Before
3839
fun setup() {
@@ -43,43 +44,62 @@ class CodeWhispererRecommendationManagerTest {
4344
runInEdtAndWait {
4445
fixture.editor.caretModel.moveToOffset(documentContentContent.length)
4546
}
47+
sut = spy(CodeWhispererRecommendationManager.getInstance())
48+
ApplicationManager.getApplication().replaceService(
49+
CodeWhispererRecommendationManager::class.java,
50+
sut,
51+
disposableRule.disposable
52+
)
4653
}
4754

4855
@Test
49-
fun `test reformatReference() should generate a new reference with span based on rangeMarker and no surfix newline char`() {
50-
// invocationOffset and markerStartOffset is of our choice as long as invocationOffset <= markerStartOffset
51-
val recommendationManager = CodeWhispererRecommendationManager()
52-
testReformatReferenceUtil(recommendationManager, documentContentSurfix = "", invocationOffset = 2, markerStartOffset = 5)
53-
testReformatReferenceUtil(recommendationManager, documentContentSurfix = "\n", invocationOffset = 2, markerStartOffset = 5)
54-
testReformatReferenceUtil(recommendationManager, documentContentSurfix = "\n\n", invocationOffset = 1, markerStartOffset = 4)
56+
fun `test overlap()`() {
57+
assertThat(sut.overlap("def", "abc")).isEqualTo("")
58+
assertThat(sut.overlap("def", "fgh")).isEqualTo("f")
59+
assertThat(sut.overlap(" ", " }")).isEqualTo(" ")
60+
assertThat(sut.overlap("abcd", "abc")).isEqualTo("")
5561
}
5662

57-
private fun testReformatReferenceUtil(
58-
recommendationManager: CodeWhispererRecommendationManager,
59-
documentContentSurfix: String,
60-
invocationOffset: Int,
61-
markerStartOffset: Int
62-
) {
63-
// insert newline characters
64-
WriteCommandAction.runWriteCommandAction(project) {
65-
fixture.editor.document.insertString(fixture.editor.caretModel.offset, documentContentSurfix)
66-
}
63+
@Test
64+
fun `test recommendation will be discarded when it's a exact match to user's input`() {
65+
val userInput = "def"
66+
val detail = sut.buildDetailContext(aRequestContext(project), userInput, listOf(aCompletion("def")), aString())
67+
assertThat(detail[0].isDiscarded).isTrue
68+
assertThat(detail[0].isTruncatedOnRight).isFalse
69+
}
6770

68-
val rangeMarker =
69-
runInEdtAndGet { fixture.editor.document.createRangeMarker(markerStartOffset, documentContentContent.length + documentContentSurfix.length) }
71+
@Test
72+
fun `test duplicated recommendation after truncation will be discarded`() {
73+
val userInput = ""
74+
sut.stub {
75+
onGeneric { findRightContextOverlap(any(), any()) } doReturn "}"
76+
onGeneric { reformatReference(any(), any()) } doReturn aCompletion("def")
77+
}
78+
val detail = sut.buildDetailContext(
79+
aRequestContext(project),
80+
userInput,
81+
listOf(aCompletion("def"), aCompletion("def}")),
82+
aString()
83+
)
84+
assertThat(detail[0].isDiscarded).isFalse
85+
assertThat(detail[0].isTruncatedOnRight).isFalse
86+
assertThat(detail[1].isDiscarded).isTrue
87+
assertThat(detail[1].isTruncatedOnRight).isTrue
88+
}
7089

71-
val reformattedReference = runInEdtAndGet { recommendationManager.reformatReference(originalReference, rangeMarker, invocationOffset) }
72-
assertThat(reformattedReference.licenseName()).isEqualTo("test_license")
73-
assertThat(reformattedReference.repository()).isEqualTo("test_repo")
74-
assertThat(reformattedReference.recommendationContentSpan().start()).isEqualTo(rangeMarker.startOffset - invocationOffset)
75-
assertThat(reformattedReference.recommendationContentSpan().end()).isEqualTo(rangeMarker.endOffset - invocationOffset - documentContentSurfix.length)
76-
val span = runInEdtAndGet {
77-
fixture.editor.document.charsSequence.subSequence(
78-
reformattedReference.recommendationContentSpan().start(),
79-
reformattedReference.recommendationContentSpan().end()
80-
)
90+
@Test
91+
fun `test blank recommendation after truncation will be discarded`() {
92+
val userInput = ""
93+
sut.stub {
94+
onGeneric { findRightContextOverlap(any(), any()) } doReturn "}"
8195
}
82-
// span should not include newline char
83-
assertThat(span.last()).isNotEqualTo('\n')
96+
val detail = sut.buildDetailContext(
97+
aRequestContext(project),
98+
userInput,
99+
listOf(aCompletion(" }")),
100+
aString()
101+
)
102+
assertThat(detail[0].isDiscarded).isTrue
103+
assertThat(detail[0].isTruncatedOnRight).isTrue
84104
}
85105
}

0 commit comments

Comments
 (0)