Merge branch 'main' into samgst/diskFullFailures

samgst-amazon · web-flow · commit 783551a45802 · 2024-11-12T10:28:13.000-08:00
diff --git a/.changes/3.39.json b/.changes/3.39.json
@@ -0,0 +1,8 @@
+{
+  "date" : "2024-11-12",
+  "version" : "3.39",
+  "entries" : [ {
+    "type" : "bugfix",
+    "description" : "Fix poor inline suggestions from Amazon Q caused by improperly formatted supplemental context"
+  } ]
+}
diff --git a/.changes/next-release/bugfix-4dc2263c-25fa-4d07-a6e1-69389f5ec001.json b/.changes/next-release/bugfix-4dc2263c-25fa-4d07-a6e1-69389f5ec001.json
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,6 @@
+# _3.39_ (2024-11-12)
+- **(Bug Fix)** Fix poor inline suggestions from Amazon Q caused by improperly formatted supplemental context
+
 # _3.38_ (2024-11-07)
 - **(Bug Fix)** Improve the position that inline chat shortcut hint is shown in editor
 - **(Bug Fix)** Improve `@workspace` index start stop strategy
diff --git a/gradle.properties b/gradle.properties
@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # Toolkit Version
-toolkitVersion=3.39-SNAPSHOT
+toolkitVersion=3.40-SNAPSHOT
 
 # Publish Settings
 publishToken=
diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/workspace/context/ProjectContextProviderTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/workspace/context/ProjectContextProviderTest.kt
@@ -364,7 +364,7 @@ class ProjectContextProviderTest {
                         .withResponseBody(
                             Body(validQueryInlineResponse)
                         )
-                        .withFixedDelay(51) // 10 sec
+                        .withFixedDelay(101) // 100 ms
                 )
             )
 
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/model/CodeWhispererModel.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/model/CodeWhispererModel.kt
@@ -15,6 +15,7 @@ import com.intellij.util.concurrency.annotations.RequiresEdt
 import software.amazon.awssdk.services.codewhispererruntime.model.Completion
 import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
+import software.aws.toolkits.jetbrains.services.amazonq.SUPPLEMENTAL_CONTEXT_TIMEOUT
 import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.PayloadContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
 import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManagerNew
@@ -73,7 +74,7 @@ data class SupplementalContextInfo(
         }
 
     val isProcessTimeout: Boolean
-        get() = latency > CodeWhispererConstants.SUPPLEMENTAL_CONTEXT_TIMEOUT
+        get() = latency > SUPPLEMENTAL_CONTEXT_TIMEOUT
 
     companion object {
         fun emptyCrossFileContextInfo(targetFileName: String): SupplementalContextInfo = SupplementalContextInfo(
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt
@@ -53,6 +53,7 @@ import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
+import software.aws.toolkits.jetbrains.services.amazonq.SUPPLEMENTAL_CONTEXT_TIMEOUT
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorManager
@@ -76,7 +77,6 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhis
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CaretMovement
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeInsightsSettingsFacade
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
-import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.SUPPLEMENTAL_CONTEXT_TIMEOUT
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getCompletionType
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getTelemetryOptOutPreference
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.notifyErrorCodeWhispererUsageLimit
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererServiceNew.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererServiceNew.kt
@@ -53,6 +53,7 @@ import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
+import software.aws.toolkits.jetbrains.services.amazonq.SUPPLEMENTAL_CONTEXT_TIMEOUT
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorManagerNew
@@ -78,7 +79,6 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhis
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CaretMovement
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeInsightsSettingsFacade
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
-import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.SUPPLEMENTAL_CONTEXT_TIMEOUT
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getCompletionType
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getTelemetryOptOutPreference
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.notifyErrorCodeWhispererUsageLimit
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt
@@ -30,7 +30,6 @@ object CodeWhispererConstants {
     const val POPUP_DELAY: Long = 250
     const val POPUP_DELAY_CHECK_INTERVAL: Long = 25
     const val IDLE_TIME_CHECK_INTERVAL: Long = 25
-    const val SUPPLEMENTAL_CONTEXT_TIMEOUT = 50L
     const val SUPPLEMETAL_CONTEXT_BUFFER = 10L
 
     val AWSTemplateKeyWordsRegex = Regex("(AWSTemplateFormatVersion|Resources|AWS::|Description)")
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererServiceTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererServiceTest.kt
@@ -19,6 +19,7 @@ import org.junit.Test
 import org.mockito.kotlin.any
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.spy
 import org.mockito.kotlin.times
@@ -35,6 +36,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.customization.Code
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
+import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.Chunk
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo
@@ -93,6 +95,33 @@ class CodeWhispererServiceTest {
         projectRule.project.replaceService(AwsConnectionManager::class.java, mock(), disposableRule.disposable)
     }
 
+    @Test
+    fun `getRequestContext should use correct fileContext and timeout to fetch supplementalContext`() = runTest {
+        val fileContextProvider = FileContextProvider.getInstance(projectRule.project)
+        val fileContextProviderSpy = spy(fileContextProvider)
+        projectRule.project.replaceService(FileContextProvider::class.java, fileContextProviderSpy, disposableRule.disposable)
+
+        val requestContext = sut.getRequestContext(
+            TriggerTypeInfo(CodewhispererTriggerType.AutoTrigger, CodeWhispererAutomatedTriggerType.Enter()),
+            editor = projectRule.fixture.editor,
+            project = projectRule.project,
+            file,
+            LatencyContext()
+        )
+
+        requestContext.awaitSupplementalContext()
+        val fileContextCaptor = argumentCaptor<FileContextInfo>()
+        verify(fileContextProviderSpy, times(1)).extractSupplementalFileContext(eq(file), fileContextCaptor.capture(), eq(100))
+        assertThat(fileContextCaptor.firstValue).isEqualTo(
+            FileContextInfo(
+                CaretContext(leftFileContext = "", rightFileContext = "public class Main {}", leftContextOnCurrentLine = ""),
+                "main.java",
+                CodeWhispererJava.INSTANCE,
+                "main.java"
+            )
+        )
+    }
+
     @Test
     fun `getRequestContext should have supplementalContext and customizatioArn if they're present`() {
         whenever(customizationConfig.activeCustomization(projectRule.project)).thenReturn(
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/Constants.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/Constants.kt
@@ -37,3 +37,5 @@ const val CODE_TRANSFORM_PREREQUISITES =
     "https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/code-transformation.html#prerequisites"
 
 const val FEATURE_EVALUATION_PRODUCT_NAME = "CodeWhisperer"
+
+const val SUPPLEMENTAL_CONTEXT_TIMEOUT = 100L
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt
@@ -27,6 +27,7 @@ import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.services.amazonq.FeatureDevSessionContext
+import software.aws.toolkits.jetbrains.services.amazonq.SUPPLEMENTAL_CONTEXT_TIMEOUT
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl
 import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
 import software.aws.toolkits.telemetry.AmazonqTelemetry
@@ -172,7 +173,7 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
         }
     }
 
-    suspend fun queryInline(query: String, filePath: String): List<InlineBm25Chunk> = withTimeout(50L) {
+    suspend fun queryInline(query: String, filePath: String): List<InlineBm25Chunk> = withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) {
         cs.async {
             val encrypted = encryptRequest(QueryInlineCompletionRequest(query, filePath))
             val r = sendMsgToLsp(LspMessage.QueryInlineCompletion, encrypted)
diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt
@@ -110,7 +110,7 @@ class DiskCache(
         clientRegistrationCache(ssoRegion).tryDeleteIfExists()
     }
 
-    override fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey): ClientRegistration? {
+    override fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey, source: String): ClientRegistration? {
         LOG.info { "loadClientRegistration for $cacheKey" }
         val cacheFile = clientRegistrationCache(cacheKey)
         val diskData = cacheFile.tryInputStreamIfExists()?.use { it.readBytes() }
@@ -144,7 +144,7 @@ class DiskCache(
                 val stage = LoadCredentialStage.ACCESS_FILE
                 AuthTelemetry.modifyConnection(
                     action = "Load cache file",
-                    source = "loadClientRegistration",
+                    source = "loadClientRegistration:$source",
                     result = Result.Failed,
                     reason = "Failed to load Client Registration",
                     reasonDesc = "Load Step:$stage failed. Cache file does not exist"
diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt
@@ -194,7 +194,7 @@ class SsoAccessTokenProvider(
 
     @Deprecated("Device authorization grant flow is deprecated")
     private fun registerDAGClient(): ClientRegistration {
-        loadDagClientRegistration()?.let {
+        loadDagClientRegistration(SourceOfLoadRegistration.REGISTER_CLIENT.toString())?.let {
             return it
         }
 
@@ -235,7 +235,7 @@ class SsoAccessTokenProvider(
     }
 
     private fun registerPkceClient(): PKCEClientRegistration {
-        loadPkceClientRegistration()?.let {
+        loadPkceClientRegistration(SourceOfLoadRegistration.REGISTER_CLIENT.toString())?.let {
             return it
         }
 
@@ -431,8 +431,8 @@ class SsoAccessTokenProvider(
         stageName = RefreshCredentialStage.LOAD_REGISTRATION
         val registration = try {
             when (currentToken) {
-                is DeviceAuthorizationGrantToken -> loadDagClientRegistration()
-                is PKCEAuthorizationGrantToken -> loadPkceClientRegistration()
+                is DeviceAuthorizationGrantToken -> loadDagClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString())
+                is PKCEAuthorizationGrantToken -> loadPkceClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString())
             }
         } catch (e: Exception) {
             val message = e.message ?: "$stageName: ${e::class.java.name}"
@@ -505,6 +505,11 @@ class SsoAccessTokenProvider(
         }
     }
 
+    enum class SourceOfLoadRegistration {
+        REGISTER_CLIENT,
+        REFRESH_TOKEN,
+    }
+
     private enum class RefreshCredentialStage {
         VALIDATE_REFRESH_TOKEN,
         LOAD_REGISTRATION,
@@ -514,13 +519,13 @@ class SsoAccessTokenProvider(
         SAVE_TOKEN,
     }
 
-    private fun loadDagClientRegistration(): ClientRegistration? =
-        cache.loadClientRegistration(dagClientRegistrationCacheKey)?.let {
+    private fun loadDagClientRegistration(source: String): ClientRegistration? =
+        cache.loadClientRegistration(dagClientRegistrationCacheKey, source)?.let {
             return it
         }
 
-    private fun loadPkceClientRegistration(): PKCEClientRegistration? =
-        cache.loadClientRegistration(pkceClientRegistrationCacheKey)?.let {
+    private fun loadPkceClientRegistration(source: String): PKCEClientRegistration? =
+        cache.loadClientRegistration(pkceClientRegistrationCacheKey, source)?.let {
             return it as PKCEClientRegistration
         }
 
diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCache.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCache.kt
@@ -7,7 +7,7 @@ interface SsoCache {
     fun invalidateClientRegistration(ssoRegion: String)
     fun invalidateAccessToken(ssoUrl: String)
 
-    fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey): ClientRegistration?
+    fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey, source: String): ClientRegistration?
     fun saveClientRegistration(cacheKey: ClientRegistrationCacheKey, registration: ClientRegistration)
     fun invalidateClientRegistration(cacheKey: ClientRegistrationCacheKey)
 
diff --git a/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCacheTest.kt b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCacheTest.kt
@@ -72,7 +72,8 @@ class DiskCacheTest {
                     startUrl = ssoUrl,
                     scopes = scopes,
                     region = ssoRegion
-                )
+                ),
+                "testSource"
             )
         ).isNull()
     }
@@ -86,7 +87,7 @@ class DiskCacheTest {
         )
         cacheLocation.resolve("223224b6f0b4702c1a984be8284fe2c9d9718759.json").writeText("badData")
 
-        assertThat(sut.loadClientRegistration(key)).isNull()
+        assertThat(sut.loadClientRegistration(key, "testSource")).isNull()
     }
 
     @Test
@@ -106,7 +107,7 @@ class DiskCacheTest {
             """.trimIndent()
         )
 
-        assertThat(sut.loadClientRegistration(key)).isNull()
+        assertThat(sut.loadClientRegistration(key, "testSource")).isNull()
     }
 
     @Test
@@ -127,7 +128,7 @@ class DiskCacheTest {
             """.trimIndent()
         )
 
-        assertThat(sut.loadClientRegistration(key)).isNull()
+        assertThat(sut.loadClientRegistration(key, "testSource")).isNull()
     }
 
     @Test
@@ -149,7 +150,7 @@ class DiskCacheTest {
             """.trimIndent()
         )
 
-        assertThat(sut.loadClientRegistration(key))
+        assertThat(sut.loadClientRegistration(key, "testSource"))
             .usingRecursiveComparison()
             .isEqualTo(
                 DeviceAuthorizationClientRegistration(
@@ -232,7 +233,7 @@ class DiskCacheTest {
                 """.trimIndent()
             )
 
-        assertThat(sut.loadClientRegistration(key))
+        assertThat(sut.loadClientRegistration(key, "testSource"))
             .usingRecursiveComparison()
             .isEqualTo(
                 PKCEClientRegistration(
@@ -338,10 +339,10 @@ class DiskCacheTest {
             )
         )
 
-        assertThat(sut.loadClientRegistration(key1))
+        assertThat(sut.loadClientRegistration(key1, "testSource"))
             .usingRecursiveComparison()
             .isEqualTo(
-                sut.loadClientRegistration(key2)
+                sut.loadClientRegistration(key2, "testSource")
             )
     }
 
@@ -365,11 +366,11 @@ class DiskCacheTest {
             region = ssoRegion
         )
 
-        assertThat(sut.loadClientRegistration(key)).isNotNull()
+        assertThat(sut.loadClientRegistration(key, "testSource")).isNotNull()
 
         sut.invalidateClientRegistration(key)
 
-        assertThat(sut.loadClientRegistration(key)).isNull()
+        assertThat(sut.loadClientRegistration(key, "testSource")).isNull()
         assertThat(cacheFile).doesNotExist()
     }
 
@@ -634,7 +635,7 @@ class DiskCacheTest {
         registration.setPosixFilePermissions(emptySet())
         assertPosixPermissions(registration, "---------")
 
-        assertThat(sut.loadClientRegistration(key)).isNotNull()
+        assertThat(sut.loadClientRegistration(key, "testSource")).isNotNull()
 
         assertPosixPermissions(registration, "rw-------")
     }
diff --git a/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProviderTest.kt b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProviderTest.kt
diff --git a/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/bearer/InteractiveBearerTokenProviderTest.kt b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/bearer/InteractiveBearerTokenProviderTest.kt

Original file line number	Diff line number	Diff line change
`@@ -364,7 +364,7 @@ class ProjectContextProviderTest {`
`364`	`364`	`.withResponseBody(`
`365`	`365`	`Body(validQueryInlineResponse)`
`366`	`366`	`)`
`367`		`- .withFixedDelay(51) // 10 sec`
	`367`	`+ .withFixedDelay(101) // 100 ms`
`368`	`368`	`)`
`369`	`369`	`)`
`370`	`370`