Skip to content

Commit 6a1739d

Browse files
Will-ShaoHuaWill-ShaoHua
authored
config(amazonq): default optin project context for internal users (#5019)
1 parent aed4ab7 commit 6a1739d

File tree

8 files changed

+70
-26
lines changed

8 files changed

+70
-26
lines changed

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/startup/AmazonQStartupActivity.kt‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,17 @@ import kotlinx.coroutines.delay
1414
import kotlinx.coroutines.time.withTimeout
1515
import software.aws.toolkits.core.utils.getLogger
1616
import software.aws.toolkits.core.utils.warn
17+
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
18+
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
19+
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
20+
import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser
1721
import software.aws.toolkits.jetbrains.core.gettingstarted.emitUserState
1822
import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextController
1923
import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindow
2024
import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindowFactory
2125
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
2226
import software.aws.toolkits.jetbrains.services.cwc.inline.InlineChatController
27+
import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
2328
import java.lang.management.ManagementFactory
2429
import java.time.Duration
2530
import java.util.concurrent.atomic.AtomicBoolean
@@ -30,6 +35,12 @@ class AmazonQStartupActivity : ProjectActivity {
3035
override suspend fun execute(project: Project) {
3136
if (ApplicationManager.getApplication().isUnitTestMode) return
3237

38+
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let {
39+
if (it is AwsBearerTokenConnection && isInternalUser(it.startUrl)) {
40+
CodeWhispererSettings.getInstance().toggleProjectContextEnabled(value = true, passive = true)
41+
}
42+
}
43+
3344
// initialize html contents in BGT so users don't have to wait when they open the tool window
3445
AmazonQToolWindow.getInstance(project)
3546
InlineChatController.getInstance(project)

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/GenerateUnitTestsAction.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,14 @@ import com.intellij.openapi.actionSystem.CommonDataKeys
99
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
1010
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
1111
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
12+
import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser
1213

1314
class GenerateUnitTestsAction : CustomAction(EditorContextCommand.GenerateUnitTests) {
1415
override fun getActionUpdateThread() = ActionUpdateThread.BGT
1516

1617
override fun update(e: AnActionEvent) {
1718
val project = e.getData(CommonDataKeys.PROJECT) ?: return
1819
val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) as? AwsBearerTokenConnection
19-
e.presentation.isEnabledAndVisible = connection?.startUrl == "https://amzn.awsapps.com/start"
20+
e.presentation.isEnabledAndVisible = isInternalUser(connection?.startUrl)
2021
}
2122
}

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,8 @@ import software.aws.toolkits.core.utils.getLogger
3636
import software.aws.toolkits.core.utils.info
3737
import software.aws.toolkits.core.utils.warn
3838
import software.aws.toolkits.jetbrains.core.coroutines.EDT
39+
import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser
3940
import software.aws.toolkits.jetbrains.services.amazonq.CHAT_IMPLICIT_PROJECT_CONTEXT_TIMEOUT
40-
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
4141
import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext
4242
import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController
4343
import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthNeededState
@@ -130,8 +130,8 @@ class ChatController private constructor(
130130
val triggerId = UUID.randomUUID().toString()
131131
var shouldAddIndexInProgressMessage: Boolean = false
132132
var shouldUseWorkspaceContext: Boolean = false
133-
val isDataCollectionGroup = CodeWhispererFeatureConfigService.getInstance().getIsDataCollectionEnabled()
134133
val startUrl = getStartUrl(context.project)
134+
val isInternalUser = isInternalUser(startUrl)
135135

136136
if (prompt.contains("@workspace")) {
137137
if (CodeWhispererSettings.getInstance().isProjectContextEnabled()) {
@@ -144,7 +144,9 @@ class ChatController private constructor(
144144
} else {
145145
sendOpenSettingsMessage(message.tabId)
146146
}
147-
} else if (CodeWhispererSettings.getInstance().isProjectContextEnabled() && isDataCollectionGroup) {
147+
} else if (CodeWhispererSettings.getInstance().isProjectContextEnabled() && isInternalUser) {
148+
// if user does not have @workspace in the prompt, but user is Amazon internal
149+
// add project context by default
148150
val projectContextController = ProjectContextController.getInstance(context.project)
149151
queryResult = projectContextController.query(prompt, timeout = CHAT_IMPLICIT_PROJECT_CONTEXT_TIMEOUT)
150152
}

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package software.aws.toolkits.jetbrains.services.cwc.controller.chat.userIntent
55

66
import software.amazon.awssdk.services.codewhispererstreaming.model.UserIntent
7+
import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser
78
import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteraction
89
import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteractionType
910
import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.FollowUpType
@@ -25,7 +26,7 @@ class UserIntentRecognizer {
2526
prompt.startsWith("Refactor") -> UserIntent.SUGGEST_ALTERNATE_IMPLEMENTATION
2627
prompt.startsWith("Fix") -> UserIntent.APPLY_COMMON_BEST_PRACTICES
2728
prompt.startsWith("Optimize") -> UserIntent.IMPROVE_CODE
28-
prompt.startsWith("Generate unit tests") && isInternalAmazonUser(startUrl) -> UserIntent.GENERATE_UNIT_TESTS
29+
prompt.startsWith("Generate unit tests") && isInternalUser(startUrl) -> UserIntent.GENERATE_UNIT_TESTS
2930
else -> null
3031
}
3132

@@ -45,6 +46,4 @@ class UserIntentRecognizer {
4546
fun getUserIntentFromOnboardingPageInteraction(interaction: OnboardingPageInteraction) = when (interaction.type) {
4647
OnboardingPageInteractionType.CwcButtonClick -> null
4748
}
48-
49-
private fun isInternalAmazonUser(startUrl: String?): Boolean = startUrl == "https://amzn.awsapps.com/start"
5049
}

‎plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,3 +146,43 @@ class CodeWhispererSettingsTest : CodeWhispererTestBase() {
146146
}
147147
}
148148
}
149+
150+
class CodeWhispererSettingUnitTest {
151+
private lateinit var sut: CodeWhispererSettings
152+
153+
@Before
154+
fun setUp() {
155+
sut = CodeWhispererSettings()
156+
sut.loadState(CodeWhispererConfiguration())
157+
}
158+
159+
@Test
160+
fun `projectContext is disabled by default`() {
161+
assertThat(sut.isProjectContextEnabled()).isFalse
162+
}
163+
164+
@Test
165+
fun `toggleProjectContext should set the value correct`() {
166+
assertThat(sut.isProjectContextEnabled()).isFalse
167+
168+
sut.toggleProjectContextEnabled(true)
169+
assertThat(sut.isProjectContextEnabled()).isTrue
170+
171+
sut.toggleProjectContextEnabled(false)
172+
assertThat(sut.isProjectContextEnabled()).isFalse
173+
}
174+
175+
@Test
176+
fun `toggleProjectContext should only set once on users behalf if passive is true`() {
177+
assertThat(sut.isProjectContextEnabled()).isFalse
178+
179+
sut.toggleProjectContextEnabled(true, passive = true)
180+
assertThat(sut.isProjectContextEnabled()).isTrue
181+
182+
sut.toggleProjectContextEnabled(false, passive = true)
183+
assertThat(sut.isProjectContextEnabled()).isTrue
184+
185+
sut.toggleProjectContextEnabled(false, passive = false)
186+
assertThat(sut.isProjectContextEnabled()).isFalse
187+
}
188+
}

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -109,9 +109,6 @@ class CodeWhispererFeatureConfigService {
109109
// 6) Add a test case for this feature.
110110
fun getTestFeature(): String = getFeatureValueForKey(TEST_FEATURE_NAME).stringValue()
111111

112-
fun getIsDataCollectionEnabled(): Boolean =
113-
getFeatureValueForKey(DATA_COLLECTION_FEATURE).stringValue() == "data-collection"
114-
115112
fun getCustomizationArnOverride(): String = getFeatureValueForKey(CUSTOMIZATION_ARN_OVERRIDE_NAME).stringValue()
116113

117114
fun getNewAutoTriggerUX(): Boolean = getFeatureValueForKey(NEW_AUTO_TRIGGER_UX).stringValue() == "TREATMENT"
@@ -131,7 +128,6 @@ class CodeWhispererFeatureConfigService {
131128
fun getInstance(): CodeWhispererFeatureConfigService = service()
132129
private const val TEST_FEATURE_NAME = "testFeature"
133130
private const val INLINE_COMPLETION = "ProjectContextV2"
134-
private const val DATA_COLLECTION_FEATURE = "IDEProjectContextDataCollection"
135131
const val CUSTOMIZATION_ARN_OVERRIDE_NAME = "customizationArnOverride"
136132
private const val NEW_AUTO_TRIGGER_UX = "newAutoTriggerUX"
137133
private val LOG = getLogger<CodeWhispererFeatureConfigService>()

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/settings/CodeWhispererSettings.kt‎

Lines changed: 8 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ import com.intellij.openapi.components.State
1111
import com.intellij.openapi.components.Storage
1212
import com.intellij.openapi.components.service
1313
import com.intellij.util.xmlb.annotations.Property
14-
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
1514

1615
@Service
1716
@State(name = "codewhispererSettings", storages = [Storage("aws.xml", roamingType = RoamingType.DISABLED)])
@@ -45,25 +44,19 @@ class CodeWhispererSettings : PersistentStateComponent<CodeWhispererConfiguratio
4544
true
4645
)
4746

48-
fun toggleProjectContextEnabled(value: Boolean) {
49-
state.value[CodeWhispererConfigurationType.IsProjectContextEnabled] = value
50-
}
51-
52-
fun isProjectContextEnabled() = getIsProjectContextEnabled()
53-
54-
private fun getIsProjectContextEnabled(): Boolean {
55-
val value = state.value.getOrDefault(CodeWhispererConfigurationType.IsProjectContextEnabled, false)
56-
val isDataCollectionGroup = CodeWhispererFeatureConfigService.getInstance().getIsDataCollectionEnabled()
57-
if (!value) {
58-
if (isDataCollectionGroup && !hasEnabledProjectContextOnce()) {
59-
toggleProjectContextEnabled(true)
47+
fun toggleProjectContextEnabled(value: Boolean, passive: Boolean = false) {
48+
if (passive) {
49+
if (!hasEnabledProjectContextOnce()) {
6050
toggleEnabledProjectContextOnce(true)
61-
return true
51+
state.value[CodeWhispererConfigurationType.IsProjectContextEnabled] = value
6252
}
53+
} else {
54+
state.value[CodeWhispererConfigurationType.IsProjectContextEnabled] = value
6355
}
64-
return value
6556
}
6657

58+
fun isProjectContextEnabled() = state.value.getOrDefault(CodeWhispererConfigurationType.IsProjectContextEnabled, false)
59+
6760
private fun hasEnabledProjectContextOnce() = state.value.getOrDefault(CodeWhispererConfigurationType.HasEnabledProjectContextOnce, false)
6861

6962
private fun toggleEnabledProjectContextOnce(value: Boolean) {

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sono/SonoConstants.kt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,3 +34,5 @@ fun ToolkitConnection?.isSono() = if (this == null) {
3434
} else {
3535
this is AwsBearerTokenConnection && this.startUrl == SONO_URL
3636
}
37+
38+
fun isInternalUser(startUrl: String?): Boolean = startUrl == "https://amzn.awsapps.com/start"

0 commit comments

Comments
 (0)