Skip to content

Commit aa08719

Browse files
aws-toolkit-automationaws-toolkit-automation
authored
Merge main into feature/dev-execution
2 parents 622df43 + 90d34be commit aa08719

File tree

87 files changed

+2074
-171
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

87 files changed

+2074
-171
lines changed

‎.changes/next-release/feature-5ac20013-1054-4be2-a5a1-84622a84f6c6.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "feature",
3+
"description" : "Enterprise users can choose their preferred Amazon Q profile to improve personalization and workflow across different business regions"
4+
}

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/clients/AmazonQCodeGenerateClient.kt‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,10 @@ import software.aws.toolkits.core.utils.error
3131
import software.aws.toolkits.core.utils.getLogger
3232
import software.aws.toolkits.core.utils.info
3333
import software.aws.toolkits.jetbrains.common.session.Intent
34-
import software.aws.toolkits.jetbrains.core.awsClient
3534
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
3635
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
3736
import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient
37+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
3838
import software.aws.toolkits.jetbrains.services.amazonqDoc.FEATURE_EVALUATION_PRODUCT_NAME
3939
import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTotalLatency
4040
import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata
@@ -72,7 +72,7 @@ class AmazonQCodeGenerateClient(private val project: Project) {
7272
fun connection() = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
7373
?: error("Attempted to use connection while one does not exist")
7474

75-
fun bearerClient() = connection().getConnectionSettings().awsClient<CodeWhispererRuntimeClient>()
75+
fun bearerClient() = QRegionProfileManager.getInstance().getQClient<CodeWhispererRuntimeClient>(project)
7676

7777
private val amazonQStreamingClient
7878
get() = AmazonQStreamingClient.getInstance(project)
@@ -88,6 +88,7 @@ class AmazonQCodeGenerateClient(private val project: Project) {
8888
}
8989
requestBuilder.optOutPreference(getTelemetryOptOutPreference())
9090
requestBuilder.userContext(docUserContext)
91+
requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn)
9192
}
9293

9394
fun sendDocMetricData(operationName: String, result: String): SendTelemetryEventResponse =
@@ -118,7 +119,9 @@ class AmazonQCodeGenerateClient(private val project: Project) {
118119
}
119120

120121
fun createTaskAssistConversation(): CreateTaskAssistConversationResponse = bearerClient().createTaskAssistConversation(
121-
CreateTaskAssistConversationRequest.builder().build()
122+
CreateTaskAssistConversationRequest.builder()
123+
.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn)
124+
.build()
122125
)
123126

124127
fun createTaskAssistUploadUrl(conversationId: String, contentChecksumSha256: String, contentLength: Long): CreateUploadUrlResponse =
@@ -137,6 +140,7 @@ class AmazonQCodeGenerateClient(private val project: Project) {
137140
)
138141
.build()
139142
)
143+
.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn)
140144
}
141145

142146
fun startTaskAssistCodeGeneration(conversationId: String, uploadId: String, userMessage: String, intent: Intent): StartTaskAssistCodeGenerationResponse =
@@ -155,6 +159,7 @@ class AmazonQCodeGenerateClient(private val project: Project) {
155159
.uploadId(uploadId)
156160
}
157161
.intent(intent.name)
162+
.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn)
158163
}
159164

160165
fun getTaskAssistCodeGeneration(conversationId: String, codeGenerationId: String): GetTaskAssistCodeGenerationResponse = bearerClient()

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

Lines changed: 55 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ package software.aws.toolkits.jetbrains.services.amazonq
66
import com.intellij.openapi.Disposable
77
import com.intellij.openapi.actionSystem.AnActionEvent
88
import com.intellij.openapi.actionSystem.DataContext
9+
import com.intellij.openapi.application.ApplicationManager
910
import com.intellij.openapi.application.runInEdt
1011
import com.intellij.openapi.components.Service
1112
import com.intellij.openapi.components.service
@@ -33,11 +34,16 @@ import software.aws.toolkits.jetbrains.core.webview.BrowserState
3334
import software.aws.toolkits.jetbrains.core.webview.LoginBrowser
3435
import software.aws.toolkits.jetbrains.core.webview.WebviewResourceHandlerFactory
3536
import software.aws.toolkits.jetbrains.isDeveloperMode
37+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QProfileSwitchIntent
38+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
39+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
3640
import software.aws.toolkits.jetbrains.services.amazonq.util.createBrowser
3741
import software.aws.toolkits.jetbrains.utils.isQConnected
3842
import software.aws.toolkits.jetbrains.utils.isQExpired
3943
import software.aws.toolkits.jetbrains.utils.isQWebviewsAvailable
4044
import software.aws.toolkits.telemetry.FeatureId
45+
import software.aws.toolkits.telemetry.MetricResult
46+
import software.aws.toolkits.telemetry.Telemetry
4147
import software.aws.toolkits.telemetry.UiTelemetry
4248
import software.aws.toolkits.telemetry.WebviewTelemetry
4349
import java.awt.event.ActionListener
@@ -204,6 +210,18 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
204210
UiTelemetry.click(project, signInOption)
205211
}
206212
}
213+
214+
is BrowserMessage.SwitchProfile -> {
215+
QRegionProfileManager.getInstance().switchProfile(
216+
project,
217+
QRegionProfile(profileName = message.profileName, arn = message.arn),
218+
intent = QProfileSwitchIntent.Auth
219+
)
220+
}
221+
222+
is BrowserMessage.PublishWebviewTelemetry -> {
223+
publishTelemetry(message)
224+
}
207225
}
208226
}
209227

@@ -244,13 +262,35 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
244262
}
245263

246264
// TODO: pass "REAUTH" if connection expires
247-
val stage = if (isQExpired(project)) {
248-
"REAUTH"
249-
} else {
250-
"START"
251-
}
265+
// Perform the potentially blocking AWS call outside the EDT to fetch available region profiles.
266+
ApplicationManager.getApplication().executeOnPooledThread {
267+
var errorMessage: String? = null
268+
val profiles: List<QRegionProfile> = try {
269+
QRegionProfileManager.getInstance().listRegionProfiles(project).orEmpty()
270+
} catch (e: Exception) {
271+
errorMessage = e.message
272+
LOG.warn { "Failed to call listRegionProfiles API" }
273+
val qConn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
274+
Telemetry.amazonq.didSelectProfile.use { span ->
275+
span.source(QProfileSwitchIntent.Auth.value)
276+
.amazonQProfileRegion(QRegionProfileManager.getInstance().activeProfile(project)?.region ?: "not-set")
277+
.ssoRegion((qConn as? AwsBearerTokenConnection)?.region)
278+
.credentialStartUrl((qConn as? AwsBearerTokenConnection)?.startUrl)
279+
.result(MetricResult.Failed)
280+
.reason(e.message)
281+
}
282+
emptyList()
283+
}
284+
285+
val stage = if (isQExpired(project)) {
286+
"REAUTH"
287+
} else if (isQConnected(project) && QRegionProfileManager.getInstance().isPendingProfileSelection(project)) {
288+
"PROFILE_SELECT"
289+
} else {
290+
"START"
291+
}
252292

253-
val jsonData = """
293+
val jsonData = """
254294
{
255295
stage: '$stage',
256296
regions: $regions,
@@ -261,10 +301,16 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
261301
},
262302
cancellable: ${state.browserCancellable},
263303
feature: '${state.feature}',
264-
existConnections: ${writeValueAsString(selectionSettings.values.map { it.currentSelection }.toList())}
304+
existConnections: ${writeValueAsString(selectionSettings.values.map { it.currentSelection }.toList())},
305+
profiles: ${writeValueAsString(profiles)},
306+
errorMessage: ${errorMessage?.let { "\"$it\"" } ?: "null"}
265307
}
266-
""".trimIndent()
267-
executeJS("window.ideClient.prepareUi($jsonData)")
308+
""".trimIndent()
309+
310+
runInEdt {
311+
executeJS("window.ideClient.prepareUi($jsonData)")
312+
}
313+
}
268314
}
269315

270316
override fun loginIAM(profileName: String, accessKey: String, secretKey: String) {

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
1919
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
2020
import software.aws.toolkits.jetbrains.core.gettingstarted.emitUserState
2121
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
22+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
2223
import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextController
2324
import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindow
2425
import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindowFactory
@@ -52,6 +53,9 @@ class AmazonQStartupActivity : ProjectActivity {
5253
CodeWhispererExplorerActionManager.getInstance().setIsFirstRestartAfterQInstall(false)
5354
}
5455
}
56+
57+
QRegionProfileManager.getInstance().validateProfile(project)
58+
5559
startLsp(project)
5660
if (runOnce.get()) return
5761
emitUserState(project)

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindow.kt‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage
2424
import software.aws.toolkits.jetbrains.services.amazonq.messages.MessageConnector
2525
import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteraction
2626
import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteractionType
27+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
2728
import software.aws.toolkits.jetbrains.services.amazonq.util.highlightCommand
2829
import software.aws.toolkits.jetbrains.services.amazonq.webview.BrowserConnector
2930
import software.aws.toolkits.jetbrains.services.amazonq.webview.FqnWebviewAdapter
@@ -127,7 +128,8 @@ class AmazonQToolWindow private constructor(
127128
isCodeScanAvailable = isCodeScanAvailable(project),
128129
isCodeTestAvailable = isCodeTestAvailable(project),
129130
isDocAvailable = isDocAvailable(project),
130-
highlightCommand = highlightCommand()
131+
highlightCommand = highlightCommand(),
132+
activeProfile = QRegionProfileManager.getInstance().takeIf { it.shouldDisplayProfileInfo(project) }?.activeProfile(project)
131133
)
132134

133135
scope.launch {

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindowFactory.kt‎

Lines changed: 31 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,10 @@ import com.intellij.ui.components.panels.Wrapper
1414
import com.intellij.util.ui.components.BorderLayoutPanel
1515
import software.aws.toolkits.core.utils.debug
1616
import software.aws.toolkits.core.utils.getLogger
17-
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
1817
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
1918
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
2019
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
2120
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
22-
import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES
2321
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenAuthState
2422
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
2523
import software.aws.toolkits.jetbrains.core.notifications.NotificationPanel
@@ -28,6 +26,9 @@ import software.aws.toolkits.jetbrains.core.webview.BrowserState
2826
import software.aws.toolkits.jetbrains.services.amazonq.QWebviewPanel
2927
import software.aws.toolkits.jetbrains.services.amazonq.RefreshQChatPanelButtonPressedListener
3028
import software.aws.toolkits.jetbrains.services.amazonq.gettingstarted.openMeetQPage
29+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
30+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
31+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
3132
import software.aws.toolkits.jetbrains.utils.isQConnected
3233
import software.aws.toolkits.jetbrains.utils.isQExpired
3334
import software.aws.toolkits.jetbrains.utils.isQWebviewsAvailable
@@ -62,7 +63,11 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware {
6263
ToolkitConnectionManagerListener.TOPIC,
6364
object : ToolkitConnectionManagerListener {
6465
override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
65-
onConnectionChanged(project, newConnection, qPanel)
66+
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let { qConn ->
67+
openMeetQPage(project)
68+
QRegionProfileManager.getInstance().validateProfile(project)
69+
}
70+
prepareChatContent(project, qPanel)
6671
}
6772
}
6873
)
@@ -71,9 +76,7 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware {
7176
RefreshQChatPanelButtonPressedListener.TOPIC,
7277
object : RefreshQChatPanelButtonPressedListener {
7378
override fun onRefresh() {
74-
runInEdt {
75-
prepareChatContent(project, qPanel)
76-
}
79+
prepareChatContent(project, qPanel)
7780
}
7881
}
7982
)
@@ -83,16 +86,23 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware {
8386
object : BearerTokenProviderListener {
8487
override fun onChange(providerId: String, newScopes: List<String>?) {
8588
if (ToolkitConnectionManager.getInstance(project).connectionStateForFeature(QConnection.getInstance()) == BearerTokenAuthState.AUTHORIZED) {
86-
val qComponent = AmazonQToolWindow.getInstance(project).component
87-
88-
runInEdt {
89-
qPanel.setContent(qComponent)
90-
}
89+
prepareChatContent(project, qPanel)
9190
}
9291
}
9392
}
9493
)
9594

95+
project.messageBus.connect(toolWindow.disposable).subscribe(
96+
QRegionProfileSelectedListener.TOPIC,
97+
object : QRegionProfileSelectedListener {
98+
override fun onProfileSelected(project: Project, profile: QRegionProfile?) {
99+
if (project.isDisposed) return
100+
AmazonQToolWindow.getInstance(project).disposeAndRecreate()
101+
prepareChatContent(project, qPanel)
102+
}
103+
}
104+
)
105+
96106
prepareChatContent(project, qPanel)
97107

98108
val content = contentManager.factory.createContent(mainPanel, null, false).also {
@@ -107,13 +117,21 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware {
107117
project: Project,
108118
qPanel: Wrapper,
109119
) {
110-
val component = if (isQConnected(project) && !isQExpired(project)) {
120+
/**
121+
* only render Q Chat when
122+
* 1. There is a Q connection
123+
* 2. Q connection is not expired
124+
* 3. User is not pending region profile selection
125+
*/
126+
val component = if (isQConnected(project) && !isQExpired(project) && !QRegionProfileManager.getInstance().isPendingProfileSelection(project)) {
111127
AmazonQToolWindow.getInstance(project).component
112128
} else {
113129
QWebviewPanel.getInstance(project).browser?.prepareBrowser(BrowserState(FeatureId.AmazonQ))
114130
QWebviewPanel.getInstance(project).component
115131
}
116-
qPanel.setContent(component)
132+
runInEdt {
133+
qPanel.setContent(component)
134+
}
117135
}
118136

119137
override fun init(toolWindow: ToolWindow) {
@@ -134,36 +152,6 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware {
134152

135153
override fun shouldBeAvailable(project: Project): Boolean = isQWebviewsAvailable()
136154

137-
private fun onConnectionChanged(project: Project, newConnection: ToolkitConnection?, qPanel: Wrapper) {
138-
val isNewConnectionForQ = newConnection?.let {
139-
(it as? AwsBearerTokenConnection)?.let { conn ->
140-
val scopeShouldHave = Q_SCOPES
141-
142-
LOG.debug { "newConnection: ${conn.id}; scope: ${conn.scopes}; scope must-have: $scopeShouldHave" }
143-
144-
scopeShouldHave.all { s -> s in conn.scopes }
145-
} ?: false
146-
} ?: false
147-
148-
if (isNewConnectionForQ) {
149-
openMeetQPage(project)
150-
}
151-
152-
QWebviewPanel.getInstance(project).browser?.prepareBrowser(BrowserState(FeatureId.AmazonQ))
153-
154-
// isQConnected alone is not robust and there is race condition (read/update connection states)
155-
val component = if (isNewConnectionForQ || (isQConnected(project) && !isQExpired(project))) {
156-
LOG.debug { "returning Q-chat window; isQConnection=$isNewConnectionForQ; hasPinnedConnection=$isNewConnectionForQ" }
157-
AmazonQToolWindow.getInstance(project).component
158-
} else {
159-
LOG.debug { "returning login window; no Q connection found" }
160-
QWebviewPanel.getInstance(project).component
161-
}
162-
runInEdt {
163-
qPanel.setContent(component)
164-
}
165-
}
166-
167155
companion object {
168156
private val LOG = getLogger<AmazonQToolWindowFactory>()
169157
const val WINDOW_ID = AMAZON_Q_WINDOW_ID

0 commit comments

Comments
 (0)