Cwspr crossfile all

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererC.kt

@@ -17,6 +17,8 @@ class CodeWhispererC private constructor() : CodeWhispererProgrammingLanguage()
 
     override fun isAutoFileScanSupported(): Boolean = true
 
+    override fun isSupplementalContextSupported() = true
+
     companion object {
         const val ID = "c"
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererCpp.kt

@@ -17,6 +17,8 @@ class CodeWhispererCpp private constructor() : CodeWhispererProgrammingLanguage(
 
     override fun isAutoFileScanSupported(): Boolean = true
 
+    override fun isSupplementalContextSupported() = true
+
     companion object {
         const val ID = "cpp"
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererCsharp.kt

@@ -6,7 +6,7 @@
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
 import software.aws.toolkits.telemetry.CodewhispererLanguage
 
 class CodeWhispererCsharp private constructor() : CodeWhispererProgrammingLanguage() {
     override val languageId: String = ID
 
     override fun toTelemetryType(): CodewhispererLanguage = CodewhispererLanguage.Csharp
@@ -17,7 +17,9 @@
 
     override fun isAutoFileScanSupported(): Boolean = true
 
+    override fun isSupplementalContextSupported() = true
+
     companion object {
         const val ID = "csharp"
 
         val INSTANCE = CodeWhispererCsharp()

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererGo.kt

@@ -17,6 +17,8 @@ class CodeWhispererGo private constructor() : CodeWhispererProgrammingLanguage()
 
     override fun isAutoFileScanSupported(): Boolean = true
 
+    override fun isSupplementalContextSupported() = true
+
     companion object {
         const val ID = "go"
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererKotlin.kt

@@ -13,6 +13,8 @@ class CodeWhispererKotlin private constructor() : CodeWhispererProgrammingLangua
 
     override fun isCodeCompletionSupported(): Boolean = true
 
+    override fun isSupplementalContextSupported() = true
+
     companion object {
         const val ID = "kotlin"
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererPhp.kt

@@ -17,6 +17,8 @@ class CodeWhispererPhp private constructor() : CodeWhispererProgrammingLanguage(
 
     override fun isAutoFileScanSupported(): Boolean = true
 
+    override fun isSupplementalContextSupported() = true
+
     companion object {
         const val ID = "php"
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererRuby.kt

@@ -17,6 +17,8 @@ class CodeWhispererRuby private constructor() : CodeWhispererProgrammingLanguage
 
     override fun isAutoFileScanSupported(): Boolean = true
 
+    override fun isSupplementalContextSupported() = true
+
     companion object {
         const val ID = "ruby"
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererRust.kt

@@ -13,6 +13,8 @@ class CodeWhispererRust private constructor() : CodeWhispererProgrammingLanguage
 
     override fun isCodeCompletionSupported(): Boolean = true
 
+    override fun isSupplementalContextSupported() = true
+
     companion object {
         const val ID = "rust"
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererScala.kt

@@ -13,6 +13,8 @@ class CodeWhispererScala private constructor() : CodeWhispererProgrammingLanguag
 
     override fun isCodeCompletionSupported(): Boolean = true
 
+    override fun isSupplementalContextSupported() = true
+
     companion object {
         const val ID = "scala"
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererShell.kt

@@ -13,6 +13,8 @@ class CodeWhispererShell private constructor() : CodeWhispererProgrammingLanguag
 
     override fun isCodeCompletionSupported(): Boolean = true
 
+    override fun isSupplementalContextSupported() = true
+
     companion object {
         const val ID = "shell"
 

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

@@ -82,6 +82,19 @@
 
     fun getCustomizationArnOverride(): String = getFeatureValueForKey(CUSTOMIZATION_ARN_OVERRIDE_NAME).stringValue()
 
+    fun getCrissfileConfig(): String {
+        // TODO: use key instead of random
+//        val key = getFeatureValueForKey(CROSSFILE_KEY).stringValue()
+        val randomNum = Math.random()
+        val group = if (randomNum < 1 / 2.0) {
+            "control"
+        } else {
+            "experiment"
+        }
+
+        return group
+    }
+
     // Get the feature value for the given key.
     // In case of a misconfiguration, it will return a default feature value of Boolean true.
     private fun getFeatureValueForKey(name: String): FeatureValue =
@@ -92,6 +105,7 @@
         fun getInstance(): CodeWhispererFeatureConfigService = service()
         private const val TEST_FEATURE_NAME = "testFeature"
         const val CUSTOMIZATION_ARN_OVERRIDE_NAME = "customizationArnOverride"
+        const val CROSSFILE_KEY = "crossfile"
         private val LOG = getLogger<CodeWhispererFeatureConfigService>()
 
         // TODO: add real feature later

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

@@ -21,7 +21,9 @@ import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiFile
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.intellij.util.messages.Topic
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.TimeoutCancellationException
+import kotlinx.coroutines.async
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
@@ -43,7 +45,6 @@ import software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingExce
 import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
-import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope
 import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
@@ -91,7 +92,7 @@ import software.aws.toolkits.telemetry.CodewhispererTriggerType
 import java.util.concurrent.TimeUnit
 
 @Service
-class CodeWhispererService : Disposable {
+class CodeWhispererService(private val coroutineScope: CoroutineScope) : Disposable {
     private val codeInsightSettingsFacade = CodeInsightsSettingsFacade()
     private var refreshFailure: Int = 0
 
@@ -193,7 +194,6 @@ class CodeWhispererService : Disposable {
         // from CodeWhispererPopupManager.cancelPopup() and CodeWhispererPopupManager.closePopup().
         // It's possible and ok that coroutine will keep running until the next time we check it's state.
         // As long as we don't show to the user extra info we are good.
-        val coroutineScope = disposableCoroutineScope(popup)
 
         var states: InvocationContext? = null
         var lastRecommendationIndex = -1
@@ -624,27 +624,29 @@ class CodeWhispererService : Disposable {
         val startFetchingTimestamp = System.currentTimeMillis()
         val isTstFile = FileContextProvider.getInstance(project).isTestFile(psiFile)
         val supplementalContext = runBlocking {
-            try {
-                withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) {
-                    FileContextProvider.getInstance(project).extractSupplementalFileContext(psiFile, fileContext)
-                }
-            } catch (e: Exception) {
-                if (e is TimeoutCancellationException) {
-                    LOG.debug {
-                        "Supplemental context fetch timed out in ${System.currentTimeMillis() - startFetchingTimestamp}ms"
+            coroutineScope.async {
+                try {
+                    withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) {
+                        FileContextProvider.getInstance(project).extractSupplementalFileContext(psiFile, fileContext)
+                    }
+                } catch (e: Exception) {
+                    if (e is TimeoutCancellationException) {
+                        LOG.debug {
+                            "Supplemental context fetch timed out in ${System.currentTimeMillis() - startFetchingTimestamp}ms"
+                        }
+                        SupplementalContextInfo(
+                            isUtg = isTstFile,
+                            contents = emptyList(),
+                            latency = System.currentTimeMillis() - startFetchingTimestamp,
+                            targetFileName = fileContext.filename,
+                            strategy = if (isTstFile) UtgStrategy.Empty else CrossFileStrategy.Empty
+                        )
+                    } else {
+                        LOG.debug { "Run into unexpected error when fetching supplemental context, error: ${e.message}" }
+                        null
                     }
-                    SupplementalContextInfo(
-                        isUtg = isTstFile,
-                        contents = emptyList(),
-                        latency = System.currentTimeMillis() - startFetchingTimestamp,
-                        targetFileName = fileContext.filename,
-                        strategy = if (isTstFile) UtgStrategy.Empty else CrossFileStrategy.Empty
-                    )
-                } else {
-                    LOG.debug { "Run into unexpected error when fetching supplemental context, error: ${e.message}" }
-                    null
                 }
-            }
+            }.await()
         }
 
         // 3. caret position

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

@@ -19,6 +19,7 @@ import kotlin.reflect.KClass
 /**
  * Component controlling codewhisperer user group settings
  */
+@Deprecated("use CodeWhispererFeatureConfigService instead")
 @Service
 @State(name = "codewhispererUserGroupSettings", storages = [Storage("aws.xml", roamingType = RoamingType.DISABLED)])
 class CodeWhispererUserGroupSettings : PersistentStateComponent<CodeWhispererUserGroupStates> {

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

@@ -9,6 +9,7 @@ import com.intellij.ui.JBColor
 import software.amazon.awssdk.regions.Region
 import software.amazon.awssdk.services.codewhispererruntime.model.AccessDeniedException
 import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException
+import software.aws.toolkits.jetbrains.isDeveloperMode
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
 import software.aws.toolkits.telemetry.CodewhispererGettingStartedTask
 import java.awt.Font
@@ -133,9 +134,14 @@ object CodeWhispererConstants {
         }
     }
     object CrossFile {
-        const val CHUNK_SIZE = 60
-        const val NUMBER_OF_LINE_IN_CHUNK = 10
+        val CHUNK_SIZE
+            get() = if (isDeveloperMode()) 200 else 60
+        val NUMBER_OF_LINE_IN_CHUNK
+            get() = if (isDeveloperMode()) 50 else 10
+
+        // TODO: 3 -> 10 when service side CR is done
         const val NUMBER_OF_CHUNK_TO_FETCH = 3
+        const val NEIGHBOR_FILES_DISTANCE = 1
     }
 
     object Utg {

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

@@ -20,18 +20,30 @@ import software.aws.toolkits.core.utils.debug
 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.isDeveloperMode
 import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererC
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCpp
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCsharp
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererGo
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJavaScript
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJsx
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererKotlin
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPhp
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererRuby
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererRust
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererScala
+import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererShell
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTsx
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTypeScript
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmingLanguage
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.Chunk
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.SupplementalContextInfo
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererFeatureConfigService
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings
 import java.io.DataInput
@@ -114,7 +126,8 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
         val language = targetContext.programmingLanguage
         val group = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
 
-        val supplementalContext = if (isTst) {
+        // if utg is not supported, use crossfile context as fallback
+        val supplementalContext = if (isTst && language.isUTGSupported()) {
             when (shouldFetchUtgContext(language, group)) {
                 true -> extractSupplementalFileContextForTst(psiFile, targetContext)
                 false -> SupplementalContextInfo.emptyUtgFileContextInfo(targetContext.filename)
@@ -174,13 +187,19 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
                 chunks.addAll(file.toCodeChunk(relativePath))
                 hasUsed.add(file)
                 if (chunks.size > CodeWhispererConstants.CrossFile.CHUNK_SIZE) {
-                    LOG.debug { "finish fetching ${CodeWhispererConstants.CrossFile.CHUNK_SIZE} chunks in ${System.currentTimeMillis() - parseFilesStart} ms" }
+                    LOG.debug {
+                        "finish fetching ${CodeWhispererConstants.CrossFile.CHUNK_SIZE} " +
+                            "chunks in ${System.currentTimeMillis() - parseFilesStart} ms from files ${hasUsed.map { it.name }}"
+                    }
                     return chunks.take(CodeWhispererConstants.CrossFile.CHUNK_SIZE)
                 }
             }
         }
 
-        LOG.debug { "finish fetching ${CodeWhispererConstants.CrossFile.CHUNK_SIZE} chunks in ${System.currentTimeMillis() - parseFilesStart} ms" }
+        LOG.debug {
+            "finish fetching ${CodeWhispererConstants.CrossFile.CHUNK_SIZE} " +
+                "chunks in ${System.currentTimeMillis() - parseFilesStart} ms from files ${hasUsed.map { it.name }}"
+        }
         return chunks.take(CodeWhispererConstants.CrossFile.CHUNK_SIZE)
     }
 
@@ -197,7 +216,18 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
 
         // step 1: prepare data
         val first60Chunks: List<Chunk> = try {
-            runReadAction { codewhispererCodeChunksIndex.getFileData(psiFile) }
+            runReadAction {
+                // for dev purpose only, not using cache when it's dev mode to monitor performance
+                if (isDeveloperMode()) {
+                    runBlocking {
+                        val fileCrawler = psiFile.programmingLanguage().fileCrawler
+                        val fileProducers = listOf<suspend (PsiFile) -> List<VirtualFile>> { psiFile -> fileCrawler.listCrossFileCandidate(psiFile) }
+                        FileContextProvider.getInstance(psiFile.project).extractCodeChunksFromFiles(psiFile, fileProducers)
+                    }
+                } else {
+                    codewhispererCodeChunksIndex.getFileData(psiFile)
+                }
+            }
         } catch (e: TimeoutCancellationException) {
             throw e
         }
@@ -299,6 +329,7 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
             }
         }
 
+        @Suppress("UNUSED_PARAMETER")
         fun shouldFetchCrossfileContext(language: CodeWhispererProgrammingLanguage, userGroup: CodeWhispererUserGroup): Boolean? {
             if (!language.isSupplementalContextSupported()) {
                 return null
@@ -312,7 +343,18 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
                 is CodeWhispererJsx,
                 is CodeWhispererTsx -> true
 
-                else -> userGroup == CodeWhispererUserGroup.CrossFile
+                is CodeWhispererC,
+                is CodeWhispererGo,
+                is CodeWhispererPhp,
+                is CodeWhispererRust,
+                is CodeWhispererKotlin,
+                is CodeWhispererCpp,
+                is CodeWhispererCsharp,
+                is CodeWhispererRuby,
+                is CodeWhispererShell,
+                is CodeWhispererScala -> CodeWhispererFeatureConfigService.getInstance().getCrissfileConfig() == "experiment"
+
+                else -> false
             }
         }
     }