make queryChat suspend and called within withTimeout

Author: Will-ShaoHua

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

@@ -214,8 +214,10 @@ class ProjectContextProviderTest {
     }
 
     @Test
-    fun `query should send correct encrypted request to lsp`() {
-        sut.query("foo")
+    fun `query should send correct encrypted request to lsp`() = runTest {
+        sut = ProjectContextProvider(project, encoderServer, this)
+        val r = sut.query("foo")
+        advanceUntilIdle()
 
         val request = QueryChatRequest("foo")
         val requestJson = mapper.writeValueAsString(request)
@@ -275,7 +277,9 @@ class ProjectContextProviderTest {
 
     @Test
     fun `query chat should return deserialized relevantDocument`() = runTest {
+        sut = ProjectContextProvider(project, encoderServer, this)
         val r = sut.query("foo")
+        advanceUntilIdle()
         assertThat(r).hasSize(2)
         assertThat(r[0]).isEqualTo(
             RelevantDocument(
@@ -364,7 +368,7 @@ class ProjectContextProviderTest {
                         .withResponseBody(
                             Body(validQueryInlineResponse)
                         )
-                        .withFixedDelay(51) // 10 sec
+                        .withFixedDelay(51)
                 )
             )
 
@@ -377,6 +381,29 @@ class ProjectContextProviderTest {
         }
     }
 
+    @Test
+    fun `queryChat should throw if time elapsed is greather than 500ms`() = runTest {
+        assertThrows<TimeoutCancellationException> {
+            sut = ProjectContextProvider(project, encoderServer, this)
+            stubFor(
+                any(urlPathEqualTo("/query")).willReturn(
+                    aResponse()
+                        .withStatus(200)
+                        .withResponseBody(
+                            Body(validQueryChatResponse)
+                        )
+                        .withFixedDelay(501)
+                )
+            )
+
+            withContext(getCoroutineBgContext()) {
+                sut.query("foo")
+            }
+
+            advanceUntilIdle()
+        }
+    }
+
     @Test
     fun `test index payload is encrypted`() = runTest {
         whenever(encoderServer.port).thenReturn(3000)
@@ -390,12 +417,9 @@ class ProjectContextProviderTest {
 
     @Test
     fun `test query payload is encrypted`() = runTest {
-        whenever(encoderServer.port).thenReturn(3000)
-        try {
-            sut.query("what does this project do")
-        } catch (e: ConnectException) {
-            // no-op
-        }
+        sut = ProjectContextProvider(project, encoderServer, this)
+        sut.query("what does this project do")
+        advanceUntilIdle()
         verify(encoderServer, times(1)).encrypt(any())
     }
 

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextController.kt

@@ -47,7 +47,7 @@ class ProjectContextController(private val project: Project, private val cs: Cor
 
     fun getProjectContextIndexComplete() = projectContextProvider.isIndexComplete.get()
 
-    fun query(prompt: String): List<RelevantDocument> {
+    suspend fun query(prompt: String): List<RelevantDocument> {
         try {
             return projectContextProvider.query(prompt)
         } catch (e: Exception) {

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt

@@ -23,7 +23,6 @@ import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.withTimeout
 import kotlinx.coroutines.yield
 import software.aws.toolkits.core.utils.debug
-import software.aws.toolkits.core.utils.error
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.core.utils.warn
@@ -159,17 +158,15 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
     }
 
     // TODO: rename queryChat
-    fun query(prompt: String): List<RelevantDocument> {
-        val encrypted = encryptRequest(QueryChatRequest(prompt))
-        val response = sendMsgToLsp(LspMessage.QueryChat, encrypted)
+    // TODO: timeout window not decided
+    suspend fun query(prompt: String): List<RelevantDocument> = withTimeout(500L) {
+        cs.async {
+            val encrypted = encryptRequest(QueryChatRequest(prompt))
+            val response = sendMsgToLsp(LspMessage.QueryChat, encrypted)
 
-        return try {
             val parsedResponse = mapper.readValue<List<Chunk>>(response.responseBody)
             queryResultToRelevantDocuments(parsedResponse)
-        } catch (e: Exception) {
-            logger.error { "error parsing query response ${e.message}" }
-            throw e
-        }
+        }.await()
     }
 
     suspend fun queryInline(query: String, filePath: String): List<InlineBm25Chunk> = withTimeout(50L) {