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 index 6c6df64d414..abde58fe26e 100644 --- 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 index c511890260e..91da779ab10 100644 --- 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 index d5c7d770106..f11c5014028 100644 --- 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 index 2e36bbce349..7f43e3d0882 100644 --- 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 index 1e11c245b10..eba9ce1641d 100644 --- 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 index 4bff4965c3d..1d5646d3869 100644 --- 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() + 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 index aba815bdc78..45341767295 100644 --- 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 index 2a3f4a19e48..c5ea9fb8177 100644 --- 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 = withTimeout(50L) { + suspend fun queryInline(query: String, filePath: String): List = withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) { cs.async { val encrypted = encryptRequest(QueryInlineCompletionRequest(query, filePath)) val r = sendMsgToLsp(LspMessage.QueryInlineCompletion, encrypted)