feat(amazonq): enable workspace context once for selected user (#4771)

* enabled workspace context once for selected users

* Update plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/TelemetryHelperTest.kt

Co-authored-by: Richard Li <[email protected]>

---------

Co-authored-by: Richard Li <[email protected]>

Author: zixlin7

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/model/Requests.kt

@@ -23,7 +23,8 @@ data class ChatRequestData(
     val userIntent: UserIntent?,
     val triggerType: TriggerType,
     val customization: CodeWhispererCustomization?,
-    val relevantTextDocuments: List<RelevantDocument>
+    val relevantTextDocuments: List<RelevantDocument>,
+    val useRelevantDocuments: Boolean
 )
 
 interface CodeNames {

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/v1/ChatSessionV1.kt

@@ -198,7 +198,7 @@ class ChatSessionV1(
      */
     private fun ChatRequestData.toChatRequest(): GenerateAssistantResponseRequest {
         val userInputMessageContextBuilder = UserInputMessageContext.builder()
-        userInputMessageContextBuilder.editorState(activeFileContext.toEditorState(relevantTextDocuments))
+        userInputMessageContextBuilder.editorState(activeFileContext.toEditorState(relevantTextDocuments, useRelevantDocuments))
         val userInputMessageContext = userInputMessageContextBuilder.build()
 
         val userInput = UserInputMessage.builder()
@@ -217,7 +217,7 @@ class ChatSessionV1(
             .build()
     }
 
-    private fun ActiveFileContext.toEditorState(relevantDocuments: List<RelevantDocument>): EditorState {
+    private fun ActiveFileContext.toEditorState(relevantDocuments: List<RelevantDocument>, useRelevantDocuments: Boolean): EditorState {
         val editorStateBuilder = EditorState.builder()
         if (fileContext != null) {
             val cursorStateBuilder = CursorState.builder()
@@ -284,6 +284,7 @@ class ChatSessionV1(
         }
 
         editorStateBuilder.relevantDocuments(documents)
+        editorStateBuilder.useRelevantDocuments(useRelevantDocuments)
         return editorStateBuilder.build()
     }
 

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt

@@ -42,6 +42,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthNeededState
 import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteraction
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteractionType
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererFeatureConfigService
 import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererConfigurable
 import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings
 import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererUserModificationTracker
@@ -127,16 +128,22 @@ class ChatController private constructor(
         var queryResult: List<RelevantDocument> = emptyList()
         val triggerId = UUID.randomUUID().toString()
         var shouldAddIndexInProgressMessage: Boolean = false
+        var shouldUseWorkspaceContext: Boolean = false
+        val isDataCollectionGroup = CodeWhispererFeatureConfigService.getInstance().getIsDataCollectionEnabled()
         if (prompt.contains("@workspace")) {
             if (CodeWhispererSettings.getInstance().isProjectContextEnabled()) {
-                val projectContextController = ProjectContextController.getInstance(context.project)
+                shouldUseWorkspaceContext = true
                 prompt = prompt.replace("@workspace", "")
+                val projectContextController = ProjectContextController.getInstance(context.project)
                 queryResult = projectContextController.query(prompt)
                 if (!projectContextController.getProjectContextIndexComplete()) shouldAddIndexInProgressMessage = true
                 logger.info { "project context relevant document count: ${queryResult.size}" }
             } else {
                 sendOpenSettingsMessage(message.tabId)
             }
+        } else if (CodeWhispererSettings.getInstance().isProjectContextEnabled() && isDataCollectionGroup) {
+            val projectContextController = ProjectContextController.getInstance(context.project)
+            queryResult = projectContextController.query(prompt)
         }
 
         handleChat(
@@ -147,7 +154,8 @@ class ChatController private constructor(
             userIntent = intentRecognizer.getUserIntentFromPromptChatMessage(message.chatMessage),
             TriggerType.Click,
             projectContextQueryResult = queryResult,
-            shouldAddIndexInProgressMessage = shouldAddIndexInProgressMessage
+            shouldAddIndexInProgressMessage = shouldAddIndexInProgressMessage,
+            shouldUseWorkspaceContext = shouldUseWorkspaceContext
         )
     }
 
@@ -378,7 +386,8 @@ class ChatController private constructor(
         userIntent: UserIntent?,
         triggerType: TriggerType,
         projectContextQueryResult: List<RelevantDocument>,
-        shouldAddIndexInProgressMessage: Boolean? = false
+        shouldAddIndexInProgressMessage: Boolean = false,
+        shouldUseWorkspaceContext: Boolean = false
     ) {
         val credentialState = authController.getAuthNeededStates(context.project).chat
         if (credentialState != null) {
@@ -397,7 +406,8 @@ class ChatController private constructor(
             userIntent = userIntent,
             triggerType = triggerType,
             customization = CodeWhispererModelConfigurator.getInstance().activeCustomization(context.project),
-            relevantTextDocuments = projectContextQueryResult
+            relevantTextDocuments = projectContextQueryResult,
+            useRelevantDocuments = shouldUseWorkspaceContext,
         )
 
         val sessionInfo = getSessionInfo(tabId)
@@ -409,7 +419,7 @@ class ChatController private constructor(
 
         // Send the request to the API and publish the responses back to the UI.
         // This is launched in a scope attached to the sessionInfo so that the Job can be cancelled on a per-session basis.
-        ChatPromptHandler(telemetryHelper).handle(tabId, triggerId, requestData, sessionInfo, shouldAddIndexInProgressMessage ?: false)
+        ChatPromptHandler(telemetryHelper).handle(tabId, triggerId, requestData, sessionInfo, shouldAddIndexInProgressMessage)
             .catch { handleError(tabId, it) }
             .onEach { context.messagesFromAppToUi.publish(it) }
             .launchIn(sessionInfo.scope)

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/chat/telemetry/TelemetryHelper.kt

@@ -88,7 +88,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
             cwsprChatHasCodeSnippet = data.activeFileContext.focusAreaContext?.codeSelection?.isNotEmpty() ?: false,
             cwsprChatProgrammingLanguage = data.activeFileContext.fileContext?.fileLanguage,
             credentialStartUrl = getStartUrl(context.project),
-            cwsprChatHasProjectContext = getIsProjectContextEnabled()
+            cwsprChatHasProjectContext = getIsProjectContextEnabled() && data.useRelevantDocuments && data.relevantTextDocuments.isNotEmpty()
         )
     }
 
@@ -116,7 +116,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
             cwsprChatConversationType = CwsprChatConversationType.Chat,
             credentialStartUrl = getStartUrl(context.project),
             codewhispererCustomizationArn = data.customization?.arn,
-            cwsprChatHasProjectContext = getIsProjectContextEnabled()
+            cwsprChatHasProjectContext = getIsProjectContextEnabled() && data.useRelevantDocuments && data.relevantTextDocuments.isNotEmpty()
         )
 
         val programmingLanguage = data.activeFileContext.fileContext?.fileLanguage
@@ -135,7 +135,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
             data.message.length,
             responseLength,
             numberOfCodeBlocks,
-            getIsProjectContextEnabled(),
+            getIsProjectContextEnabled() && data.useRelevantDocuments && data.relevantTextDocuments.isNotEmpty(),
             data.customization
         ).also {
             logger.debug {
@@ -265,7 +265,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
                     cwsprChatInteractionTarget = message.link,
                     cwsprChatHasReference = null,
                     credentialStartUrl = getStartUrl(context.project),
-                    cwsprChatHasProjectContext = getIsProjectContextEnabled(),
+                    cwsprChatHasProjectContext = getIsProjectContextEnabled()
                 )
                 ChatInteractWithMessageEvent.builder().apply {
                     conversationId(getConversationId(message.tabId).orEmpty())

plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/TelemetryHelperTest.kt

@@ -140,7 +140,8 @@ class TelemetryHelperTest {
             userIntent = UserIntent.IMPROVE_CODE,
             triggerType = TriggerType.Hotkeys,
             customization = mockCustomization,
-            relevantTextDocuments = listOf()
+            relevantTextDocuments = emptyList(),
+            useRelevantDocuments = true,
         )
         private val response = ChatMessage(
             tabId = tabId,

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

@@ -80,6 +80,8 @@ class CodeWhispererFeatureConfigService {
     // 6) Add a test case for this feature.
     fun getTestFeature(): String = getFeatureValueForKey(TEST_FEATURE_NAME).stringValue()
 
+    fun getIsDataCollectionEnabled(): Boolean = getFeatureValueForKey(DATA_COLLECTION_FEATURE).stringValue() == "data-collection"
+
     fun getCustomizationArnOverride(): String = getFeatureValueForKey(CUSTOMIZATION_ARN_OVERRIDE_NAME).stringValue()
 
     // Get the feature value for the given key.
@@ -91,6 +93,7 @@ class CodeWhispererFeatureConfigService {
     companion object {
         fun getInstance(): CodeWhispererFeatureConfigService = service()
         private const val TEST_FEATURE_NAME = "testFeature"
+        private const val DATA_COLLECTION_FEATURE = "IDEProjectContextDataCollection"
         const val CUSTOMIZATION_ARN_OVERRIDE_NAME = "customizationArnOverride"
         private val LOG = getLogger<CodeWhispererFeatureConfigService>()
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/settings/CodeWhispererSettings.kt

@@ -11,6 +11,7 @@ import com.intellij.openapi.components.State
 import com.intellij.openapi.components.Storage
 import com.intellij.openapi.components.service
 import com.intellij.util.xmlb.annotations.Property
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererFeatureConfigService
 
 @Service
 @State(name = "codewhispererSettings", storages = [Storage("aws.xml", roamingType = RoamingType.DISABLED)])
@@ -48,7 +49,26 @@ class CodeWhispererSettings : PersistentStateComponent<CodeWhispererConfiguratio
         state.value[CodeWhispererConfigurationType.IsProjectContextEnabled] = value
     }
 
-    fun isProjectContextEnabled() = state.value.getOrDefault(CodeWhispererConfigurationType.IsProjectContextEnabled, false)
+    fun isProjectContextEnabled() = getIsProjectContextEnabled()
+
+    private fun getIsProjectContextEnabled(): Boolean {
+        val value = state.value.getOrDefault(CodeWhispererConfigurationType.IsProjectContextEnabled, false)
+        val isDataCollectionGroup = CodeWhispererFeatureConfigService.getInstance().getIsDataCollectionEnabled()
+        if (!value) {
+            if (isDataCollectionGroup && !hasEnabledProjectContextOnce()) {
+                toggleProjectContextEnabled(true)
+                toggleEnabledProjectContextOnce(true)
+                return true
+            }
+        }
+        return value
+    }
+
+    private fun hasEnabledProjectContextOnce() = state.value.getOrDefault(CodeWhispererConfigurationType.HasEnabledProjectContextOnce, false)
+
+    private fun toggleEnabledProjectContextOnce(value: Boolean) {
+        state.value[CodeWhispererConfigurationType.HasEnabledProjectContextOnce] = value
+    }
 
     fun isProjectContextGpu() = state.value.getOrDefault(CodeWhispererConfigurationType.IsProjectContextGpu, false)
 
@@ -106,6 +126,7 @@ enum class CodeWhispererConfigurationType {
     IsAutoUpdateFeatureNotificationShownOnce,
     IsProjectContextEnabled,
     IsProjectContextGpu,
+    HasEnabledProjectContextOnce
 }
 
 enum class CodeWhispererIntConfigurationType {

plugins/core/sdk-codegen/codegen-resources/codewhispererstreaming/service-2.json

@@ -134,6 +134,10 @@
             "event":true,
             "sensitive":true
         },
+        "Boolean":{
+            "type":"boolean",
+            "box":true
+        },
         "ChatHistory":{
             "type":"list",
             "member":{"shape":"ChatMessage"},
@@ -276,7 +280,8 @@
             "members":{
                 "document":{"shape":"TextDocument"},
                 "cursorState":{"shape":"CursorState"},
-                "relevantDocuments": {"shape": "RelevantDocumentList"}
+                "relevantDocuments": {"shape": "RelevantDocumentList"},
+                "useRelevantDocuments": {"shape": "Boolean"}
             }
         },
         "ExportContext":{
@@ -765,4 +770,4 @@
             }
         }
     }
-}
+}