config(amazonq): default optin project context for internal users (#5019

)
aws · Nov 13, 2024 · 6a1739d · 6a1739d
1 parent aed4ab7
commit 6a1739d
Show file tree

Hide file tree

Showing 8 changed files with 70 additions and 26 deletions.
diff --git a/...ty/src/software/aws/toolkits/jetbrains/services/amazonq/startup/AmazonQStartupActivity.kt b/...ty/src/software/aws/toolkits/jetbrains/services/amazonq/startup/AmazonQStartupActivity.kt
@@ -14,12 +14,17 @@ import kotlinx.coroutines.delay
 import kotlinx.coroutines.time.withTimeout
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.warn
+import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
+import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
+import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
+import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser
 import software.aws.toolkits.jetbrains.core.gettingstarted.emitUserState
 import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextController
 import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindow
 import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindowFactory
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
 import software.aws.toolkits.jetbrains.services.cwc.inline.InlineChatController
+import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
 import java.lang.management.ManagementFactory
 import java.time.Duration
 import java.util.concurrent.atomic.AtomicBoolean
@@ -30,6 +35,12 @@ class AmazonQStartupActivity : ProjectActivity {
     override suspend fun execute(project: Project) {
         if (ApplicationManager.getApplication().isUnitTestMode) return
 
+        ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let {
+            if (it is AwsBearerTokenConnection && isInternalUser(it.startUrl)) {
+                CodeWhispererSettings.getInstance().toggleProjectContextEnabled(value = true, passive = true)
+            }
+        }
+
         // initialize html contents in BGT so users don't have to wait when they open the tool window
         AmazonQToolWindow.getInstance(project)
         InlineChatController.getInstance(project)

diff --git a/...nity/src/software/aws/toolkits/jetbrains/services/cwc/commands/GenerateUnitTestsAction.kt b/...nity/src/software/aws/toolkits/jetbrains/services/cwc/commands/GenerateUnitTestsAction.kt
@@ -9,13 +9,14 @@ import com.intellij.openapi.actionSystem.CommonDataKeys
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
+import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser
 
 class GenerateUnitTestsAction : CustomAction(EditorContextCommand.GenerateUnitTests) {
     override fun getActionUpdateThread() = ActionUpdateThread.BGT
 
     override fun update(e: AnActionEvent) {
         val project = e.getData(CommonDataKeys.PROJECT) ?: return
         val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) as? AwsBearerTokenConnection
-        e.presentation.isEnabledAndVisible = connection?.startUrl == "https://amzn.awsapps.com/start"
+        e.presentation.isEnabledAndVisible = isInternalUser(connection?.startUrl)
     }
 }
diff --git a/...s-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt b/...s-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt
@@ -36,8 +36,8 @@ 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.core.coroutines.EDT
+import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser
 import software.aws.toolkits.jetbrains.services.amazonq.CHAT_IMPLICIT_PROJECT_CONTEXT_TIMEOUT
-import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
 import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext
 import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController
 import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthNeededState
@@ -130,8 +130,8 @@ class ChatController private constructor(
         val triggerId = UUID.randomUUID().toString()
         var shouldAddIndexInProgressMessage: Boolean = false
         var shouldUseWorkspaceContext: Boolean = false
-        val isDataCollectionGroup = CodeWhispererFeatureConfigService.getInstance().getIsDataCollectionEnabled()
         val startUrl = getStartUrl(context.project)
+        val isInternalUser = isInternalUser(startUrl)
 
         if (prompt.contains("@workspace")) {
             if (CodeWhispererSettings.getInstance().isProjectContextEnabled()) {
@@ -144,7 +144,9 @@ class ChatController private constructor(
             } else {
                 sendOpenSettingsMessage(message.tabId)
             }
-        } else if (CodeWhispererSettings.getInstance().isProjectContextEnabled() && isDataCollectionGroup) {
+        } else if (CodeWhispererSettings.getInstance().isProjectContextEnabled() && isInternalUser) {
+            // if user does not have @workspace in the prompt, but user is Amazon internal
+            // add project context by default
             val projectContextController = ProjectContextController.getInstance(context.project)
             queryResult = projectContextController.query(prompt, timeout = CHAT_IMPLICIT_PROJECT_CONTEXT_TIMEOUT)
         }

diff --git a/...re/aws/toolkits/jetbrains/services/cwc/controller/chat/userIntent/UserIntentRecognizer.kt b/...re/aws/toolkits/jetbrains/services/cwc/controller/chat/userIntent/UserIntentRecognizer.kt
@@ -4,6 +4,7 @@
 package software.aws.toolkits.jetbrains.services.cwc.controller.chat.userIntent
 
 import software.amazon.awssdk.services.codewhispererstreaming.model.UserIntent
+import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteraction
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteractionType
 import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.FollowUpType
@@ -25,7 +26,7 @@ class UserIntentRecognizer {
         prompt.startsWith("Refactor") -> UserIntent.SUGGEST_ALTERNATE_IMPLEMENTATION
         prompt.startsWith("Fix") -> UserIntent.APPLY_COMMON_BEST_PRACTICES
         prompt.startsWith("Optimize") -> UserIntent.IMPROVE_CODE
-        prompt.startsWith("Generate unit tests") && isInternalAmazonUser(startUrl) -> UserIntent.GENERATE_UNIT_TESTS
+        prompt.startsWith("Generate unit tests") && isInternalUser(startUrl) -> UserIntent.GENERATE_UNIT_TESTS
         else -> null
     }
 
@@ -45,6 +46,4 @@ class UserIntentRecognizer {
     fun getUserIntentFromOnboardingPageInteraction(interaction: OnboardingPageInteraction) = when (interaction.type) {
         OnboardingPageInteractionType.CwcButtonClick -> null
     }
-
-    private fun isInternalAmazonUser(startUrl: String?): Boolean = startUrl == "https://amzn.awsapps.com/start"
 }
diff --git a/...y/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt b/...y/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt
@@ -146,3 +146,43 @@ class CodeWhispererSettingsTest : CodeWhispererTestBase() {
         }
     }
 }
+
+class CodeWhispererSettingUnitTest {
+    private lateinit var sut: CodeWhispererSettings
+
+    @Before
+    fun setUp() {
+        sut = CodeWhispererSettings()
+        sut.loadState(CodeWhispererConfiguration())
+    }
+
+    @Test
+    fun `projectContext is disabled by default`() {
+        assertThat(sut.isProjectContextEnabled()).isFalse
+    }
+
+    @Test
+    fun `toggleProjectContext should set the value correct`() {
+        assertThat(sut.isProjectContextEnabled()).isFalse
+
+        sut.toggleProjectContextEnabled(true)
+        assertThat(sut.isProjectContextEnabled()).isTrue
+
+        sut.toggleProjectContextEnabled(false)
+        assertThat(sut.isProjectContextEnabled()).isFalse
+    }
+
+    @Test
+    fun `toggleProjectContext should only set once on users behalf if passive is true`() {
+        assertThat(sut.isProjectContextEnabled()).isFalse
+
+        sut.toggleProjectContextEnabled(true, passive = true)
+        assertThat(sut.isProjectContextEnabled()).isTrue
+
+        sut.toggleProjectContextEnabled(false, passive = true)
+        assertThat(sut.isProjectContextEnabled()).isTrue
+
+        sut.toggleProjectContextEnabled(false, passive = false)
+        assertThat(sut.isProjectContextEnabled()).isFalse
+    }
+}
diff --git a/...src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt b/...src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt
@@ -109,9 +109,6 @@ 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()
 
     fun getNewAutoTriggerUX(): Boolean = getFeatureValueForKey(NEW_AUTO_TRIGGER_UX).stringValue() == "TREATMENT"
@@ -131,7 +128,6 @@ class CodeWhispererFeatureConfigService {
         fun getInstance(): CodeWhispererFeatureConfigService = service()
         private const val TEST_FEATURE_NAME = "testFeature"
         private const val INLINE_COMPLETION = "ProjectContextV2"
-        private const val DATA_COLLECTION_FEATURE = "IDEProjectContextDataCollection"
         const val CUSTOMIZATION_ARN_OVERRIDE_NAME = "customizationArnOverride"
         private const val NEW_AUTO_TRIGGER_UX = "newAutoTriggerUX"
         private val LOG = getLogger<CodeWhispererFeatureConfigService>()

diff --git a/...jetbrains-community/src/software/aws/toolkits/jetbrains/settings/CodeWhispererSettings.kt b/...jetbrains-community/src/software/aws/toolkits/jetbrains/settings/CodeWhispererSettings.kt
@@ -11,7 +11,6 @@ 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.amazonq.CodeWhispererFeatureConfigService
 
 @Service
 @State(name = "codewhispererSettings", storages = [Storage("aws.xml", roamingType = RoamingType.DISABLED)])
@@ -45,25 +44,19 @@ class CodeWhispererSettings : PersistentStateComponent<CodeWhispererConfiguratio
         true
     )
 
-    fun toggleProjectContextEnabled(value: Boolean) {
-        state.value[CodeWhispererConfigurationType.IsProjectContextEnabled] = value
-    }
-
-    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)
+    fun toggleProjectContextEnabled(value: Boolean, passive: Boolean = false) {
+        if (passive) {
+            if (!hasEnabledProjectContextOnce()) {
                 toggleEnabledProjectContextOnce(true)
-                return true
+                state.value[CodeWhispererConfigurationType.IsProjectContextEnabled] = value
             }
+        } else {
+            state.value[CodeWhispererConfigurationType.IsProjectContextEnabled] = value
         }
-        return value
     }
 
+    fun isProjectContextEnabled() = state.value.getOrDefault(CodeWhispererConfigurationType.IsProjectContextEnabled, false)
+
     private fun hasEnabledProjectContextOnce() = state.value.getOrDefault(CodeWhispererConfigurationType.HasEnabledProjectContextOnce, false)
 
     private fun toggleEnabledProjectContextOnce(value: Boolean) {

diff --git a/...ains-community/src/software/aws/toolkits/jetbrains/core/credentials/sono/SonoConstants.kt b/...ains-community/src/software/aws/toolkits/jetbrains/core/credentials/sono/SonoConstants.kt
@@ -34,3 +34,5 @@ fun ToolkitConnection?.isSono() = if (this == null) {
 } else {
     this is AwsBearerTokenConnection && this.startUrl == SONO_URL
 }
+
+fun isInternalUser(startUrl: String?): Boolean = startUrl == "https://amzn.awsapps.com/start"