Skip to content

Commit b342ff9

Browse files
Will-ShaoHuaWill-ShaoHua
authored
fix supplemental context being built in EDT (#4765)
1 parent dab6b29 commit b342ff9

File tree

8 files changed

+261
-106
lines changed

8 files changed

+261
-106
lines changed

‎.changes/next-release/bugfix-eb2b9d18-dfc2-46fd-b772-df33d48d81d3.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" : "Fix Q building supplemental context under EDT which might slow or block the UI"
4+
}

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt‎

Lines changed: 44 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,13 @@ import com.intellij.psi.PsiDocumentManager
2121
import com.intellij.psi.PsiFile
2222
import com.intellij.util.concurrency.annotations.RequiresEdt
2323
import com.intellij.util.messages.Topic
24-
import kotlinx.coroutines.TimeoutCancellationException
24+
import kotlinx.coroutines.CoroutineScope
25+
import kotlinx.coroutines.Deferred
26+
import kotlinx.coroutines.Job
27+
import kotlinx.coroutines.async
2528
import kotlinx.coroutines.delay
2629
import kotlinx.coroutines.isActive
2730
import kotlinx.coroutines.launch
28-
import kotlinx.coroutines.runBlocking
29-
import kotlinx.coroutines.withTimeout
3031
import software.amazon.awssdk.core.exception.SdkServiceException
3132
import software.amazon.awssdk.core.util.DefaultSdkAutoConstructList
3233
import software.amazon.awssdk.services.codewhisperer.model.CodeWhispererException
@@ -43,6 +44,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingExce
4344
import software.aws.toolkits.core.utils.debug
4445
import software.aws.toolkits.core.utils.getLogger
4546
import software.aws.toolkits.core.utils.info
47+
import software.aws.toolkits.core.utils.warn
4648
import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope
4749
import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
4850
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
@@ -76,9 +78,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhisperer
7678
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getTelemetryOptOutPreference
7779
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.notifyErrorCodeWhispererUsageLimit
7880
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.promptReAuth
79-
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CrossFileStrategy
8081
import software.aws.toolkits.jetbrains.services.codewhisperer.util.FileContextProvider
81-
import software.aws.toolkits.jetbrains.services.codewhisperer.util.UtgStrategy
8282
import software.aws.toolkits.jetbrains.utils.isInjectedText
8383
import software.aws.toolkits.jetbrains.utils.isQExpired
8484
import software.aws.toolkits.jetbrains.utils.isRunningOnCWNotSupportedRemoteBackend
@@ -91,7 +91,7 @@ import software.aws.toolkits.telemetry.CodewhispererTriggerType
9191
import java.util.concurrent.TimeUnit
9292

9393
@Service
94-
class CodeWhispererService : Disposable {
94+
class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
9595
private val codeInsightSettingsFacade = CodeInsightsSettingsFacade()
9696
private var refreshFailure: Int = 0
9797

@@ -184,7 +184,7 @@ class CodeWhispererService : Disposable {
184184
invokeCodeWhispererInBackground(requestContext)
185185
}
186186

187-
private fun invokeCodeWhispererInBackground(requestContext: RequestContext) {
187+
internal fun invokeCodeWhispererInBackground(requestContext: RequestContext): Job {
188188
val popup = CodeWhispererPopupManager.getInstance().initPopup()
189189
Disposer.register(popup) { CodeWhispererInvocationStatus.getInstance().finishInvocation() }
190190

@@ -198,15 +198,16 @@ class CodeWhispererService : Disposable {
198198
var states: InvocationContext? = null
199199
var lastRecommendationIndex = -1
200200

201-
val responseIterable = CodeWhispererClientAdaptor.getInstance(requestContext.project).generateCompletionsPaginator(
202-
buildCodeWhispererRequest(
203-
requestContext.fileContextInfo,
204-
requestContext.supplementalContext,
205-
requestContext.customizationArn
206-
)
207-
)
208-
coroutineScope.launch {
201+
val job = coroutineScope.launch {
209202
try {
203+
val responseIterable = CodeWhispererClientAdaptor.getInstance(requestContext.project).generateCompletionsPaginator(
204+
buildCodeWhispererRequest(
205+
requestContext.fileContextInfo,
206+
requestContext.awaitSupplementalContext(),
207+
requestContext.customizationArn
208+
)
209+
)
210+
210211
var startTime = System.nanoTime()
211212
requestContext.latencyContext.codewhispererPreprocessingEnd = System.nanoTime()
212213
requestContext.latencyContext.paginationAllCompletionsStart = System.nanoTime()
@@ -413,6 +414,8 @@ class CodeWhispererService : Disposable {
413414
CodeWhispererInvocationStatus.getInstance().setInvocationComplete()
414415
}
415416
}
417+
418+
return job
416419
}
417420

418421
@RequiresEdt
@@ -621,29 +624,12 @@ class CodeWhispererService : Disposable {
621624

622625
// the upper bound for supplemental context duration is 50ms
623626
// 2. supplemental context
624-
val startFetchingTimestamp = System.currentTimeMillis()
625-
val isTstFile = FileContextProvider.getInstance(project).isTestFile(psiFile)
626-
val supplementalContext = runBlocking {
627+
val supplementalContext = cs.async {
627628
try {
628-
withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) {
629-
FileContextProvider.getInstance(project).extractSupplementalFileContext(psiFile, fileContext)
630-
}
629+
FileContextProvider.getInstance(project).extractSupplementalFileContext(psiFile, fileContext, timeout = SUPPLEMENTAL_CONTEXT_TIMEOUT)
631630
} catch (e: Exception) {
632-
if (e is TimeoutCancellationException) {
633-
LOG.debug {
634-
"Supplemental context fetch timed out in ${System.currentTimeMillis() - startFetchingTimestamp}ms"
635-
}
636-
SupplementalContextInfo(
637-
isUtg = isTstFile,
638-
contents = emptyList(),
639-
latency = System.currentTimeMillis() - startFetchingTimestamp,
640-
targetFileName = fileContext.filename,
641-
strategy = if (isTstFile) UtgStrategy.Empty else CrossFileStrategy.Empty
642-
)
643-
} else {
644-
LOG.debug { "Run into unexpected error when fetching supplemental context, error: ${e.message}" }
645-
null
646-
}
631+
LOG.warn { "Run into unexpected error when fetching supplemental context, error: ${e.message}" }
632+
null
647633
}
648634
}
649635

@@ -825,11 +811,31 @@ data class RequestContext(
825811
val triggerTypeInfo: TriggerTypeInfo,
826812
val caretPosition: CaretPosition,
827813
val fileContextInfo: FileContextInfo,
828-
val supplementalContext: SupplementalContextInfo?,
814+
private val supplementalContextDeferred: Deferred<SupplementalContextInfo?>,
829815
val connection: ToolkitConnection?,
830816
val latencyContext: LatencyContext,
831817
val customizationArn: String?
832-
)
818+
) {
819+
// TODO: should make the entire getRequestContext() suspend function instead of making supplemental context only
820+
var supplementalContext: SupplementalContextInfo? = null
821+
private set
822+
get() = when (field) {
823+
null -> {
824+
if (!supplementalContextDeferred.isCompleted) {
825+
error("attempt to access supplemental context before awaiting the deferred")
826+
} else {
827+
null
828+
}
829+
}
830+
831+
else -> field
832+
}
833+
834+
suspend fun awaitSupplementalContext(): SupplementalContextInfo? {
835+
supplementalContext = supplementalContextDeferred.await()
836+
return supplementalContext
837+
}
838+
}
833839

834840
data class ResponseContext(
835841
val sessionId: String,

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

Lines changed: 56 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import com.intellij.util.gist.GistManager
1414
import com.intellij.util.io.DataExternalizer
1515
import kotlinx.coroutines.TimeoutCancellationException
1616
import kotlinx.coroutines.runBlocking
17+
import kotlinx.coroutines.withTimeout
1718
import kotlinx.coroutines.yield
1819
import org.jetbrains.annotations.VisibleForTesting
1920
import software.aws.toolkits.core.utils.debug
@@ -83,7 +84,7 @@ private object CodeWhispererCodeChunkExternalizer : DataExternalizer<List<Chunk>
8384
interface FileContextProvider {
8485
fun extractFileContext(editor: Editor, psiFile: PsiFile): FileContextInfo
8586

86-
suspend fun extractSupplementalFileContext(psiFile: PsiFile, fileContext: FileContextInfo): SupplementalContextInfo?
87+
suspend fun extractSupplementalFileContext(psiFile: PsiFile, fileContext: FileContextInfo, timeout: Long): SupplementalContextInfo?
8788

8889
suspend fun extractCodeChunksFromFiles(psiFile: PsiFile, fileProducers: List<suspend (PsiFile) -> List<VirtualFile>>): List<Chunk>
8990

@@ -108,56 +109,73 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
108109
* for focal files, e.g. "MainTest.java" -> "Main.java", "test_main.py" -> "main.py"
109110
* for the most relevant file -> we extract "keywords" from files opened in editor then get the one with the highest similarity with target file
110111
*/
111-
override suspend fun extractSupplementalFileContext(psiFile: PsiFile, targetContext: FileContextInfo): SupplementalContextInfo? {
112+
override suspend fun extractSupplementalFileContext(psiFile: PsiFile, targetContext: FileContextInfo, timeout: Long): SupplementalContextInfo? {
112113
val startFetchingTimestamp = System.currentTimeMillis()
113114
val isTst = isTestFile(psiFile)
114-
val language = targetContext.programmingLanguage
115-
val group = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
116-
117-
val supplementalContext = if (isTst) {
118-
when (shouldFetchUtgContext(language, group)) {
119-
true -> extractSupplementalFileContextForTst(psiFile, targetContext)
120-
false -> SupplementalContextInfo.emptyUtgFileContextInfo(targetContext.filename)
121-
null -> {
122-
LOG.debug { "UTG is not supporting ${targetContext.programmingLanguage.languageId}" }
123-
null
124-
}
125-
}
126-
} else {
127-
when (shouldFetchCrossfileContext(language, group)) {
128-
true -> extractSupplementalFileContextForSrc(psiFile, targetContext)
129-
false -> SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
130-
null -> {
131-
LOG.debug { "Crossfile is not supporting ${targetContext.programmingLanguage.languageId}" }
132-
null
115+
return try {
116+
withTimeout(timeout) {
117+
val language = targetContext.programmingLanguage
118+
val group = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
119+
120+
val supplementalContext = if (isTst) {
121+
when (shouldFetchUtgContext(language, group)) {
122+
true -> extractSupplementalFileContextForTst(psiFile, targetContext)
123+
false -> SupplementalContextInfo.emptyUtgFileContextInfo(targetContext.filename)
124+
null -> {
125+
LOG.debug { "UTG is not supporting ${targetContext.programmingLanguage.languageId}" }
126+
null
127+
}
128+
}
129+
} else {
130+
when (shouldFetchCrossfileContext(language, group)) {
131+
true -> extractSupplementalFileContextForSrc(psiFile, targetContext)
132+
false -> SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
133+
null -> {
134+
LOG.debug { "Crossfile is not supporting ${targetContext.programmingLanguage.languageId}" }
135+
null
136+
}
137+
}
133138
}
134-
}
135-
}
136139

137-
return supplementalContext?.let {
138-
if (it.contents.isNotEmpty()) {
139-
val logStr = buildString {
140-
append("Successfully fetched supplemental context.")
141-
it.contents.forEachIndexed { index, chunk ->
142-
append(
143-
"""
140+
return@withTimeout supplementalContext?.let {
141+
if (it.contents.isNotEmpty()) {
142+
val logStr = buildString {
143+
append("Successfully fetched supplemental context.")
144+
it.contents.forEachIndexed { index, chunk ->
145+
append(
146+
"""
144147
|
145148
| Chunk ${index + 1}:
146149
| path = ${chunk.path},
147150
| score = ${chunk.score},
148151
| contentLength = ${chunk.content.length}
149152
|
150-
""".trimMargin()
151-
)
153+
""".trimMargin()
154+
)
155+
}
156+
}
157+
158+
LOG.info { logStr }
159+
} else {
160+
LOG.warn { "Failed to fetch supplemental context, empty list." }
152161
}
153-
}
154162

155-
LOG.info { logStr }
156-
} else {
157-
LOG.warn { "Failed to fetch supplemental context, empty list." }
163+
it.copy(latency = System.currentTimeMillis() - startFetchingTimestamp)
164+
}
158165
}
159-
160-
it.copy(latency = System.currentTimeMillis() - startFetchingTimestamp)
166+
} catch (e: TimeoutCancellationException) {
167+
LOG.debug {
168+
"Supplemental context fetch timed out in ${System.currentTimeMillis() - startFetchingTimestamp}ms"
169+
}
170+
SupplementalContextInfo(
171+
isUtg = isTst,
172+
contents = emptyList(),
173+
latency = System.currentTimeMillis() - startFetchingTimestamp,
174+
targetFileName = targetContext.filename,
175+
strategy = if (isTst) UtgStrategy.Empty else CrossFileStrategy.Empty
176+
)
177+
} catch (e: Exception) {
178+
throw e
161179
}
162180
}
163181

‎plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ import com.intellij.testFramework.fixtures.CodeInsightTestFixture
1717
import com.intellij.testFramework.replaceService
1818
import com.intellij.testFramework.runInEdtAndGet
1919
import com.intellij.testFramework.runInEdtAndWait
20+
import kotlinx.coroutines.async
21+
import kotlinx.coroutines.runBlocking
2022
import org.assertj.core.api.Assertions.assertThat
2123
import org.junit.After
2224
import org.junit.Before
@@ -161,7 +163,17 @@ internal class CodeWhispererCodeCoverageTrackerTestPython : CodeWhispererCodeCov
161163
mock(),
162164
mock(),
163165
FileContextInfo(mock(), pythonFileName, CodeWhispererPython.INSTANCE),
164-
SupplementalContextInfo(isUtg = false, contents = emptyList(), targetFileName = "", strategy = CrossFileStrategy.Empty, latency = 0L),
166+
runBlocking {
167+
async {
168+
SupplementalContextInfo(
169+
isUtg = false,
170+
contents = emptyList(),
171+
targetFileName = "",
172+
strategy = CrossFileStrategy.Empty,
173+
latency = 0L
174+
)
175+
}
176+
},
165177
null,
166178
mock(),
167179
aString()

‎plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFileContextProviderTest.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -319,7 +319,7 @@ class CodeWhispererFileContextProviderTest {
319319
runReadAction {
320320
val fileContext = sut.extractFileContext(fixture.editor, psiFiles[0])
321321

322-
val supplementalContext = runBlocking { sut.extractSupplementalFileContext(psiFiles[0], fileContext) }
322+
val supplementalContext = runBlocking { sut.extractSupplementalFileContext(psiFiles[0], fileContext, timeout = 50) }
323323
assertThat(supplementalContext?.contents).isNotNull.isNotEmpty
324324
}
325325

@@ -365,7 +365,7 @@ class CodeWhispererFileContextProviderTest {
365365
runReadAction {
366366
val fileContext = aFileContextInfo(CodeWhispererJava.INSTANCE)
367367
val supplementalContext = runBlocking {
368-
sut.extractSupplementalFileContext(tstFile, fileContext)
368+
sut.extractSupplementalFileContext(tstFile, fileContext, 50)
369369
}
370370
assertThat(supplementalContext?.contents)
371371
.isNotNull

0 commit comments

Comments
 (0)