Skip to content

Commit ff27d01

Browse files
Will-ShaoHuaWill-ShaoHua
committed
Merge branch 'lsp-client-suspend-version' into lsp-client
2 parents 13e04d2 + 4da830b commit ff27d01

File tree

6 files changed

+145
-95
lines changed

6 files changed

+145
-95
lines changed

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

Lines changed: 49 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,13 @@ import com.intellij.openapi.application.ApplicationManager
1717
import com.intellij.openapi.project.Project
1818
import com.intellij.testFramework.DisposableRule
1919
import com.intellij.testFramework.replaceService
20+
import kotlinx.coroutines.ExperimentalCoroutinesApi
21+
import kotlinx.coroutines.TimeoutCancellationException
22+
import kotlinx.coroutines.test.StandardTestDispatcher
2023
import kotlinx.coroutines.test.TestScope
24+
import kotlinx.coroutines.test.advanceUntilIdle
2125
import kotlinx.coroutines.test.runTest
26+
import kotlinx.coroutines.withContext
2227
import org.assertj.core.api.Assertions.assertThat
2328
import org.junit.Before
2429
import org.junit.Rule
@@ -31,6 +36,7 @@ import org.mockito.kotlin.stub
3136
import org.mockito.kotlin.times
3237
import org.mockito.kotlin.verify
3338
import org.mockito.kotlin.whenever
39+
import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineBgContext
3440
import software.aws.toolkits.jetbrains.services.amazonq.project.EncoderServer
3541
import software.aws.toolkits.jetbrains.services.amazonq.project.IndexRequest
3642
import software.aws.toolkits.jetbrains.services.amazonq.project.IndexUpdateMode
@@ -45,9 +51,9 @@ import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
4551
import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
4652
import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
4753
import java.net.ConnectException
48-
import java.util.concurrent.TimeoutException
4954
import kotlin.test.Test
5055

56+
@OptIn(ExperimentalCoroutinesApi::class)
5157
class ProjectContextProviderTest {
5258
@Rule
5359
@JvmField
@@ -69,12 +75,14 @@ class ProjectContextProviderTest {
6975

7076
private val mapper = jacksonObjectMapper()
7177

78+
private val dispatcher = StandardTestDispatcher()
79+
7280
@Before
7381
fun setup() {
7482
encoderServer = spy(EncoderServer(project))
7583
encoderServer.stub { on { port } doReturn wireMock.port() }
7684

77-
sut = ProjectContextProvider(project, encoderServer, TestScope())
85+
sut = ProjectContextProvider(project, encoderServer, TestScope(context = dispatcher))
7886

7987
// initialization
8088
stubFor(any(urlPathEqualTo("/initialize")).willReturn(aResponse().withStatus(200).withResponseBody(Body("initialize response"))))
@@ -225,16 +233,17 @@ class ProjectContextProviderTest {
225233
}
226234

227235
@Test
228-
fun `queryInline should send correct encrypted request to lsp`() {
236+
fun `queryInline should send correct encrypted request to lsp`() = runTest {
237+
sut = ProjectContextProvider(project, encoderServer, this)
229238
sut.queryInline("foo", "Foo.java")
239+
advanceUntilIdle()
230240

231241
val request = QueryInlineCompletionRequest("foo", "Foo.java")
232242
val requestJson = mapper.writeValueAsString(request)
233243

234244
assertThat(mapper.readTree(requestJson)).isEqualTo(mapper.readTree("""{ "query": "foo", "filePath": "Foo.java" }"""))
235245

236246
val encryptedRequest = encoderServer.encrypt(requestJson)
237-
238247
wireMock.verify(
239248
1,
240249
postRequestedFor(urlPathEqualTo("/queryInlineProjectContext"))
@@ -284,27 +293,35 @@ class ProjectContextProviderTest {
284293

285294
@Test
286295
fun `query inline should throw if resultset not deserializable`() {
287-
stubFor(
288-
any(urlPathEqualTo("/queryInlineProjectContext")).willReturn(
289-
aResponse().withStatus(200).withResponseBody(
290-
Body(
291-
"""
296+
assertThrows<Exception> {
297+
runTest {
298+
sut = ProjectContextProvider(project, encoderServer, this)
299+
stubFor(
300+
any(urlPathEqualTo("/queryInlineProjectContext")).willReturn(
301+
aResponse().withStatus(200).withResponseBody(
302+
Body(
303+
"""
292304
[
293305
"foo", "bar"
294306
]
295-
""".trimIndent()
307+
""".trimIndent()
308+
)
309+
)
296310
)
297311
)
298-
)
299-
)
300312

301-
assertThrows<Exception> {
302-
sut.queryInline("foo", "filepath")
313+
assertThrows<Exception> {
314+
sut.queryInline("foo", "filepath")
315+
advanceUntilIdle()
316+
}
317+
}
303318
}
304319
}
305320

306321
@Test
307322
fun `query inline should return deserialized bm25 chunks`() = runTest {
323+
sut = ProjectContextProvider(project, encoderServer, this)
324+
advanceUntilIdle()
308325
val r = sut.queryInline("foo", "filepath")
309326
assertThat(r).hasSize(3)
310327
assertThat(r[0]).isEqualTo(
@@ -337,20 +354,26 @@ class ProjectContextProviderTest {
337354
}
338355

339356
@Test
340-
fun `queryInline should throw if time elapsed is greater than 50ms`() {
341-
stubFor(
342-
any(urlPathEqualTo("/queryInlineProjectContext")).willReturn(
343-
aResponse()
344-
.withStatus(200)
345-
.withResponseBody(
346-
Body(validQueryInlineResponse)
347-
)
348-
.withFixedDelay(51) // 10 sec
357+
fun `queryInline should throw if time elapsed is greater than 50ms`() = runTest {
358+
assertThrows<TimeoutCancellationException> {
359+
sut = ProjectContextProvider(project, encoderServer, this)
360+
stubFor(
361+
any(urlPathEqualTo("/queryInlineProjectContext")).willReturn(
362+
aResponse()
363+
.withStatus(200)
364+
.withResponseBody(
365+
Body(validQueryInlineResponse)
366+
)
367+
.withFixedDelay(51) // 10 sec
368+
)
349369
)
350-
)
351370

352-
assertThrows<TimeoutException> {
353-
sut.queryInline("foo", "bar")
371+
// it won't throw if it's executed within TestDispatcher context
372+
withContext(getCoroutineBgContext()) {
373+
sut.queryInline("foo", "bar")
374+
}
375+
376+
advanceUntilIdle()
354377
}
355378
}
356379

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ object CodeWhispererConstants {
3131
const val POPUP_DELAY_CHECK_INTERVAL: Long = 25
3232
const val IDLE_TIME_CHECK_INTERVAL: Long = 25
3333
const val SUPPLEMENTAL_CONTEXT_TIMEOUT = 50L
34+
const val SUPPLEMETAL_CONTEXT_BUFFER = 10L
3435

3536
val AWSTemplateKeyWordsRegex = Regex("(AWSTemplateFormatVersion|Resources|AWS::|Description)")
3637
val AWSTemplateCaseInsensitiveKeyWordsRegex = Regex("(cloudformation|cfn|template|description)")

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

Lines changed: 69 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import com.intellij.util.gist.GistManager
1515
import com.intellij.util.io.DataExternalizer
1616
import kotlinx.coroutines.TimeoutCancellationException
1717
import kotlinx.coroutines.async
18+
import kotlinx.coroutines.awaitAll
1819
import kotlinx.coroutines.runBlocking
1920
import kotlinx.coroutines.withContext
2021
import kotlinx.coroutines.withTimeout
@@ -38,6 +39,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmi
3839
import software.aws.toolkits.jetbrains.services.codewhisperer.model.Chunk
3940
import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo
4041
import software.aws.toolkits.jetbrains.services.codewhisperer.model.SupplementalContextInfo
42+
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.SUPPLEMETAL_CONTEXT_BUFFER
4143
import java.io.DataInput
4244
import java.io.DataOutput
4345
import java.util.Collections
@@ -117,54 +119,55 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
117119
val startFetchingTimestamp = System.currentTimeMillis()
118120
val isTst = readAction { isTestFile(psiFile) }
119121
return try {
120-
withTimeout(timeout) {
121-
val language = targetContext.programmingLanguage
122-
123-
val supplementalContext = if (isTst) {
124-
when (shouldFetchUtgContext(language)) {
125-
true -> extractSupplementalFileContextForTst(psiFile, targetContext)
126-
false -> SupplementalContextInfo.emptyUtgFileContextInfo(targetContext.filename)
127-
null -> {
128-
LOG.debug { "UTG is not supporting ${targetContext.programmingLanguage.languageId}" }
129-
null
130-
}
122+
val language = targetContext.programmingLanguage
123+
124+
val supplementalContext = if (isTst) {
125+
when (shouldFetchUtgContext(language)) {
126+
true -> withTimeout(timeout) { extractSupplementalFileContextForTst(psiFile, targetContext) }
127+
false -> SupplementalContextInfo.emptyUtgFileContextInfo(targetContext.filename)
128+
null -> {
129+
LOG.debug { "UTG is not supporting ${targetContext.programmingLanguage.languageId}" }
130+
null
131131
}
132-
} else {
133-
when (shouldFetchCrossfileContext(language)) {
134-
true -> extractSupplementalFileContextForSrc(psiFile, targetContext)
135-
false -> SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
136-
null -> {
137-
LOG.debug { "Crossfile is not supporting ${targetContext.programmingLanguage.languageId}" }
138-
null
139-
}
132+
}
133+
} else {
134+
when (shouldFetchCrossfileContext(language)) {
135+
// we need this buffer 10ms as when project context timeout by 50ms,
136+
// the entire [extractSupplementalFileContextForSrc] call will time out and not even return openTabsContext
137+
true -> withTimeout(timeout + SUPPLEMETAL_CONTEXT_BUFFER) { extractSupplementalFileContextForSrc(psiFile, targetContext) }
138+
false -> SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
139+
null -> {
140+
LOG.debug { "Crossfile is not supporting ${targetContext.programmingLanguage.languageId}" }
141+
null
140142
}
141143
}
144+
}
142145

143-
return@withTimeout supplementalContext?.let {
144-
if (it.contents.isNotEmpty()) {
145-
val logStr = buildString {
146-
append("Successfully fetched supplemental context with strategy ${it.strategy} with ${it.latency} ms")
147-
it.contents.forEachIndexed { index, chunk ->
148-
append(
149-
"""
146+
return supplementalContext?.let {
147+
if (it.contents.isNotEmpty()) {
148+
val logStr = buildString {
149+
append("Successfully fetched supplemental context with strategy ${it.strategy} with ${it.latency} ms")
150+
it.contents.forEachIndexed { index, chunk ->
151+
append(
152+
"""
150153
|
151154
| Chunk ${index + 1}:
152155
| path = ${chunk.path},
153156
| score = ${chunk.score},
154157
| contentLength = ${chunk.content.length}
155158
|
156-
""".trimMargin()
157-
)
158-
}
159+
""".trimMargin()
160+
)
159161
}
160-
161-
LOG.info { logStr }
162-
} else {
163-
LOG.warn { "Failed to fetch supplemental context, empty list." }
164162
}
165163

166-
it.copy(latency = System.currentTimeMillis() - startFetchingTimestamp)
164+
LOG.info { logStr }
165+
} else {
166+
LOG.warn { "Failed to fetch supplemental context, empty list." }
167167
}
168+
169+
// TODO: fix this latency
170+
it.copy(latency = System.currentTimeMillis() - startFetchingTimestamp)
168171
}
169172
} catch (e: TimeoutCancellationException) {
170173
LOG.debug {
@@ -215,47 +218,60 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
215218

216219
val query = generateQuery(targetContext)
217220

218-
val projectContextDeferred = if (CodeWhispererFeatureConfigService.getInstance().getInlineCompletion()) {
219-
withContext(coroutineContext) {
221+
val contexts = withContext(coroutineContext) {
222+
val projectContextDeferred1 = if (CodeWhispererFeatureConfigService.getInstance().getInlineCompletion()) {
220223
async {
221224
val t0 = System.currentTimeMillis()
222225
val r = fetchProjectContext(query, psiFile, targetContext)
223226
val t1 = System.currentTimeMillis()
224-
LOG.debug { "fetchProjectContext cost ${t1 - t0} ms" }
227+
LOG.debug {
228+
buildString {
229+
append("time elapse for fetching project context=${t1 - t0}ms; ")
230+
append("numberOfChunks=${r.contents.size}; ")
231+
append("totalLength=${r.contentLength}")
232+
}
233+
}
225234

226235
r
227236
}
237+
} else {
238+
null
228239
}
229-
} else {
230-
null
231-
}
232240

233-
val openTabsContextDeferred = withContext(coroutineContext) {
234-
async {
241+
val openTabsContextDeferred1 = async {
235242
val t0 = System.currentTimeMillis()
236243
val r = fetchOpenTabsContext(query, psiFile, targetContext)
237244
val t1 = System.currentTimeMillis()
238-
LOG.debug { "fetchOpenTabsContext cost ${t1 - t0} ms" }
245+
LOG.debug {
246+
buildString {
247+
append("time elapse for open tabs context=${t1 - t0}ms; ")
248+
append("numberOfChunks=${r.contents.size}; ")
249+
append("totalLength=${r.contentLength}")
250+
}
251+
}
239252

240253
r
241254
}
242-
}
243255

244-
if (projectContextDeferred == null) {
245-
return openTabsContextDeferred.await()
246-
} else {
247-
val projectContext = projectContextDeferred.await()
248-
return if (projectContext.contents.isNotEmpty()) {
249-
projectContext
256+
if (projectContextDeferred1 != null) {
257+
awaitAll(projectContextDeferred1, openTabsContextDeferred1)
250258
} else {
251-
val openTabsContext = openTabsContextDeferred.await()
252-
openTabsContext
259+
awaitAll(openTabsContextDeferred1)
253260
}
254261
}
262+
263+
val projectContext = contexts.find { it.strategy == CrossFileStrategy.ProjectContext }
264+
val openTabsContext = contexts.find { it.strategy == CrossFileStrategy.OpenTabsBM25 }
265+
266+
return if (projectContext != null && projectContext.contents.isNotEmpty()) {
267+
projectContext
268+
} else {
269+
openTabsContext ?: SupplementalContextInfo.emptyCrossFileContextInfo(targetContext.filename)
270+
}
255271
}
256272

257273
@VisibleForTesting
258-
fun fetchProjectContext(query: String, psiFile: PsiFile, targetContext: FileContextInfo): SupplementalContextInfo {
274+
suspend fun fetchProjectContext(query: String, psiFile: PsiFile, targetContext: FileContextInfo): SupplementalContextInfo {
259275
val response = ProjectContextController.getInstance(project).queryInline(query, psiFile.virtualFile?.path ?: "")
260276

261277
return SupplementalContextInfo(

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

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ class CodeWhispererFileContextProviderTest {
138138

139139
@Test
140140
fun `should call both and use openTabsContext if projectContext is empty when it's enabled`() = runTest {
141-
mockProjectContext.stub { on { queryInline(any(), any()) } doReturn emptyList() }
141+
mockProjectContext.stub { onBlocking { queryInline(any(), any()) }.doReturn(emptyList()) }
142142
featureConfigService.stub { on { getInlineCompletion() } doReturn true }
143143
sut = spy(sut)
144144

@@ -158,7 +158,7 @@ class CodeWhispererFileContextProviderTest {
158158

159159
// move to projectContextControllerTest
160160
@Test
161-
fun `projectContextController should return empty result if provider throws`() {
161+
fun `projectContextController should return empty result if provider throws`() = runTest {
162162
mockConstruction(ProjectContextProvider::class.java).use { providerContext ->
163163
mockConstruction(EncoderServer::class.java).use { serverContext ->
164164
assertThat(providerContext.constructed()).hasSize(0)
@@ -178,11 +178,15 @@ class CodeWhispererFileContextProviderTest {
178178
@Test
179179
fun `should use project context if it is present`() = runTest {
180180
mockProjectContext.stub {
181-
on { queryInline(any(), any()) } doReturn listOf(
182-
InlineBm25Chunk("project_context1", "path1", 0.0),
183-
InlineBm25Chunk("project_context2", "path2", 0.0),
184-
InlineBm25Chunk("project_context3", "path3", 0.0),
185-
)
181+
runBlocking {
182+
doReturn(
183+
listOf(
184+
InlineBm25Chunk("project_context1", "path1", 0.0),
185+
InlineBm25Chunk("project_context2", "path2", 0.0),
186+
InlineBm25Chunk("project_context3", "path3", 0.0),
187+
)
188+
).whenever(it).queryInline(any(), any())
189+
}
186190
}
187191
featureConfigService.stub { on { getInlineCompletion() } doReturn true }
188192
sut = spy(sut)

0 commit comments

Comments
 (0)