Skip to content

Commit cb08c9a

Browse files
andrewyuqandrewyuq
authored
CodeWhisperer: Fix indexOutOfBoundException in reformat (#3960)
* Fix indexOutOfBoundException in reformat 1. Remove adjustLineIndent since it's almost no longer needed, indents from model will be trusted. 2. Keep the reformatReference since reference offset still needs to be updated after right context truncation. 3. Add tests * Add change log
1 parent 1f5e736 commit cb08c9a

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)