Skip to content

Commit 1faaa9e

Browse files
aws-toolkit-automationaws-toolkit-automation
authored
Merge main into feature/q-lsp
2 parents d03a641 + 1a57043 commit 1faaa9e

File tree

14 files changed

+378
-219
lines changed

14 files changed

+378
-219
lines changed

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

Lines changed: 80 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,10 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
220220
)
221221
}
222222

223+
is BrowserMessage.ListProfiles -> {
224+
handleListProfilesMessage()
225+
}
226+
223227
is BrowserMessage.PublishWebviewTelemetry -> {
224228
publishTelemetry(message)
225229
}
@@ -262,60 +266,41 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
262266
writeValueAsString(it)
263267
}
264268

265-
// TODO: pass "REAUTH" if connection expires
266-
// Perform the potentially blocking AWS call outside the EDT to fetch available region profiles.
267-
ApplicationManager.getApplication().executeOnPooledThread {
268-
val stage = if (isQExpired(project)) {
269-
"REAUTH"
270-
} else if (isQConnected(project) && QRegionProfileManager.getInstance().isPendingProfileSelection(project)) {
271-
"PROFILE_SELECT"
272-
} else {
273-
"START"
274-
}
275-
276-
var errorMessage: String? = null
277-
var profiles: List<QRegionProfile> = emptyList()
269+
val stage = if (isQExpired(project)) {
270+
"REAUTH"
271+
} else if (isQConnected(project) && QRegionProfileManager.getInstance().isPendingProfileSelection(project)) {
272+
"PROFILE_SELECT"
273+
} else {
274+
"START"
275+
}
278276

279-
if (stage == "PROFILE_SELECT") {
280-
try {
281-
profiles = QRegionProfileManager.getInstance().listRegionProfiles(project).orEmpty()
282-
if (profiles.size == 1) {
283-
LOG.debug { "User only have access to 1 Q profile, auto-selecting profile ${profiles.first().profileName} for ${project.name}" }
284-
QRegionProfileManager.getInstance().switchProfile(project, profiles.first(), QProfileSwitchIntent.Update)
285-
}
286-
} catch (e: Exception) {
287-
errorMessage = e.message
288-
LOG.warn { "Failed to call listRegionProfiles API" }
289-
val qConn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
290-
Telemetry.amazonq.didSelectProfile.use { span ->
291-
span.source(QProfileSwitchIntent.Auth.value)
292-
.amazonQProfileRegion(QRegionProfileManager.getInstance().activeProfile(project)?.region ?: "not-set")
293-
.ssoRegion((qConn as? AwsBearerTokenConnection)?.region)
294-
.credentialStartUrl((qConn as? AwsBearerTokenConnection)?.startUrl)
295-
.result(MetricResult.Failed)
296-
.reason(e.message)
277+
when (stage) {
278+
"PROFILE_SELECT" -> {
279+
val jsonData = """
280+
{
281+
stage: '$stage',
282+
status: 'pending'
297283
}
298-
}
284+
""".trimIndent()
285+
executeJS("window.ideClient.prepareUi($jsonData)")
299286
}
300287

301-
val jsonData = """
302-
{
303-
stage: '$stage',
304-
regions: $regions,
305-
idcInfo: {
306-
profileName: '${lastLoginIdcInfo.profileName}',
307-
startUrl: '${lastLoginIdcInfo.startUrl}',
308-
region: '${lastLoginIdcInfo.region}'
309-
},
310-
cancellable: ${state.browserCancellable},
311-
feature: '${state.feature}',
312-
existConnections: ${writeValueAsString(selectionSettings.values.map { it.currentSelection }.toList())},
313-
profiles: ${writeValueAsString(profiles)},
314-
errorMessage: ${errorMessage?.let { "\"$it\"" } ?: "null"}
315-
}
316-
""".trimIndent()
288+
else -> {
289+
val jsonData = """
290+
{
291+
stage: '$stage',
292+
regions: $regions,
293+
idcInfo: {
294+
profileName: '${lastLoginIdcInfo.profileName}',
295+
startUrl: '${lastLoginIdcInfo.startUrl}',
296+
region: '${lastLoginIdcInfo.region}'
297+
},
298+
cancellable: ${state.browserCancellable},
299+
feature: '${state.feature}',
300+
existConnections: ${writeValueAsString(selectionSettings.values.map { it.currentSelection }.toList())},
301+
}
302+
""".trimIndent()
317303

318-
runInEdt {
319304
executeJS("window.ideClient.prepareUi($jsonData)")
320305
}
321306
}
@@ -330,6 +315,52 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
330315
jcefBrowser.loadHTML(getWebviewHTML(webScriptUri, query))
331316
}
332317

318+
private fun handleListProfilesMessage() {
319+
ApplicationManager.getApplication().executeOnPooledThread {
320+
var errorMessage = ""
321+
val profiles = try {
322+
QRegionProfileManager.getInstance().listRegionProfiles(project)
323+
} catch (e: Exception) {
324+
e.message?.let {
325+
errorMessage = it
326+
}
327+
LOG.warn { "Failed to call listRegionProfiles API: $errorMessage" }
328+
val qConn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
329+
Telemetry.amazonq.didSelectProfile.use { span ->
330+
span.source(QProfileSwitchIntent.Auth.value)
331+
.amazonQProfileRegion(QRegionProfileManager.getInstance().activeProfile(project)?.region ?: "not-set")
332+
.ssoRegion((qConn as? AwsBearerTokenConnection)?.region)
333+
.credentialStartUrl((qConn as? AwsBearerTokenConnection)?.startUrl)
334+
.result(MetricResult.Failed)
335+
.reason(e.message)
336+
}
337+
338+
null
339+
}
340+
341+
// auto-select the profile if users only have 1 and don't show the UI
342+
if (profiles?.size == 1) {
343+
LOG.debug { "User only have access to 1 Q profile, auto-selecting profile ${profiles.first().profileName} for ${project.name}" }
344+
QRegionProfileManager.getInstance().switchProfile(project, profiles.first(), QProfileSwitchIntent.Update)
345+
return@executeOnPooledThread
346+
}
347+
348+
// required EDT as this entire block is executed on thread pool
349+
runInEdt {
350+
val jsonData = """
351+
{
352+
stage: 'PROFILE_SELECT',
353+
status: '${if (profiles != null) "succeeded" else "failed"}',
354+
profiles: ${writeValueAsString(profiles ?: "")},
355+
errorMessage: '$errorMessage'
356+
}
357+
""".trimIndent()
358+
359+
executeJS("window.ideClient.prepareUi($jsonData)")
360+
}
361+
}
362+
}
363+
333364
companion object {
334365
private val LOG = getLogger<QWebviewBrowser>()
335366
private const val WEB_SCRIPT_URI = "http://webview/js/getStart.js"

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,8 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware {
8585
object : BearerTokenProviderListener {
8686
override fun onChange(providerId: String, newScopes: List<String>?) {
8787
if (ToolkitConnectionManager.getInstance(project).connectionStateForFeature(QConnection.getInstance()) == BearerTokenAuthState.AUTHORIZED) {
88-
prepareChatContent(project, qPanel)
88+
AmazonQToolWindow.getInstance(project).disposeAndRecreate()
89+
qPanel.setContent(AmazonQToolWindow.getInstance(project).component)
8990
}
9091
}
9192
}

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

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ package software.aws.toolkits.jetbrains.services.codewhisperer
66
import com.intellij.openapi.project.Project
77
import com.intellij.testFramework.DisposableRule
88
import com.intellij.testFramework.ProjectRule
9+
import com.intellij.testFramework.replaceService
910
import com.intellij.util.xmlb.XmlSerializer
1011
import org.assertj.core.api.Assertions.assertThat
1112
import org.jdom.output.XMLOutputter
@@ -15,7 +16,9 @@ import org.junit.Test
1516
import org.mockito.kotlin.any
1617
import org.mockito.kotlin.doReturn
1718
import org.mockito.kotlin.mock
19+
import org.mockito.kotlin.spy
1820
import org.mockito.kotlin.stub
21+
import org.mockito.kotlin.whenever
1922
import software.amazon.awssdk.core.pagination.sync.SdkIterable
2023
import software.amazon.awssdk.regions.Region
2124
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
@@ -34,6 +37,7 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
3437
import software.aws.toolkits.jetbrains.core.credentials.logoutFromSsoConnection
3538
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
3639
import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES
40+
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenAuthState
3741
import software.aws.toolkits.jetbrains.core.region.MockRegionProviderRule
3842
import software.aws.toolkits.jetbrains.services.amazonq.profile.QEndpoints
3943
import software.aws.toolkits.jetbrains.services.amazonq.profile.QProfileResources
@@ -85,6 +89,10 @@ class QRegionProfileManagerTest {
8589
sut = QRegionProfileManager()
8690
val conn = authRule.createConnection(ManagedSsoProfile(ssoRegion = "us-east-1", startUrl = "", scopes = Q_SCOPES))
8791
ToolkitConnectionManager.getInstance(project).switchConnection(conn)
92+
val realManager = ToolkitConnectionManager.getInstance(project)
93+
val managerSpy = spy(realManager)
94+
doReturn(BearerTokenAuthState.AUTHORIZED).whenever(managerSpy).connectionStateForFeature(QConnection.getInstance())
95+
project.replaceService(ToolkitConnectionManager::class.java, managerSpy, disposableRule.disposable)
8896
}
8997

9098
@Test
@@ -106,7 +114,7 @@ class QRegionProfileManagerTest {
106114
logoutFromSsoConnection(project, it)
107115
}
108116
}
109-
117+
ToolkitConnectionManager.getInstance(project).switchConnection(null)
110118
assertThat(ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())).isNull()
111119
assertThat(sut.activeProfile(project)).isNull()
112120
}
@@ -384,4 +392,28 @@ class QRegionProfileManagerTest {
384392
val profileList = actualState.connectionIdToProfileList["conn-123"]
385393
assertThat(profileList).isEqualTo(2)
386394
}
395+
396+
@Test
397+
fun `getIdcConnectionOrNull handles NOT_AUTH and AUTHORIZED correctly`() {
398+
val managerSpy = ToolkitConnectionManager.getInstance(project)
399+
doReturn(BearerTokenAuthState.NOT_AUTHENTICATED).whenever(managerSpy)
400+
.connectionStateForFeature(QConnection.getInstance())
401+
402+
// NOT AUTHORIZED
403+
val notAuthConn = sut.getIdcConnectionOrNull(project)
404+
assertThat(notAuthConn).isNull()
405+
406+
doReturn(BearerTokenAuthState.AUTHORIZED)
407+
.whenever(managerSpy).connectionStateForFeature(QConnection.getInstance())
408+
409+
// AUTHORIZED
410+
val normalConn = authRule.createConnection(
411+
ManagedSsoProfile(ssoRegion = "us-east-1", startUrl = "", scopes = Q_SCOPES)
412+
)
413+
managerSpy.switchConnection(normalConn)
414+
415+
val normalConnectionResult = sut.getIdcConnectionOrNull(project)
416+
assertThat(normalConnectionResult).isNotNull()
417+
assertThat(normalConnectionResult).isEqualTo(normalConn)
418+
}
387419
}

‎plugins/amazonq/shared/jetbrains-community/resources/software/aws/toolkits/resources/AmazonQBundle.properties

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
action.q.hello.description=Hello description
21
amazonqInlineChat.hint.edit = Edit
32
amazonqInlineChat.popup.accept=Accept \u23CE
43
amazonqInlineChat.popup.cancel=Cancel \u238B
@@ -10,9 +9,9 @@ amazonqInlineChat.popup.title=Enter Instructions for Q
109
amazonq.refresh.panel=Refresh Chat Session
1110
amazonq.title=Amazon Q
1211
amazonq.workspace.settings.open.prompt=Workspace index is now enabled. You can disable it from Amazon Q settings.
13-
action.q.profile.usage.text=You changed profile
14-
action.q.profile.usage=You're using the '<b>{0}</b>' profile for Amazon Q.
15-
action.q.switchProfiles.text=Change profile
12+
action.q.profile.usage.text=You changed your profile
13+
action.q.profile.usage=You''re using the ''<b>{0}</b>'' profile for Amazon Q.
14+
action.q.switchProfiles.text=Change Profile
1615
action.q.switchProfiles.dialog.text=Amazon Q Developer Profile
1716
action.q.switchProfiles.dialog.account.label=Account: {0}
1817
action.q.switchProfiles.dialog.panel.text=Change your Q Developer profile

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

Lines changed: 4 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -7,17 +7,10 @@ import com.intellij.icons.AllIcons
77
import com.intellij.openapi.actionSystem.ActionUpdateThread
88
import com.intellij.openapi.actionSystem.AnAction
99
import com.intellij.openapi.actionSystem.AnActionEvent
10-
import com.intellij.openapi.application.ApplicationManager
1110
import com.intellij.openapi.project.DumbAware
12-
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
13-
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
14-
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
15-
import software.aws.toolkits.jetbrains.services.amazonq.profile.QProfileSwitchIntent
1611
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileDialog
1712
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
1813
import software.aws.toolkits.resources.AmazonQBundle.message
19-
import software.aws.toolkits.telemetry.MetricResult
20-
import software.aws.toolkits.telemetry.Telemetry
2114

2215
class QSwitchProfilesAction : AnAction(message("action.q.switchProfiles.text")), DumbAware {
2316

@@ -29,30 +22,9 @@ class QSwitchProfilesAction : AnAction(message("action.q.switchProfiles.text")),
2922

3023
override fun actionPerformed(e: AnActionEvent) {
3124
val project = e.project ?: return
32-
ApplicationManager.getApplication().executeOnPooledThread {
33-
val profiles = try {
34-
QRegionProfileManager.getInstance().listRegionProfiles(project)
35-
} catch (e: Exception) {
36-
val conn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) as? AwsBearerTokenConnection
37-
Telemetry.amazonq.didSelectProfile.use { span ->
38-
span.source(QProfileSwitchIntent.User.value)
39-
.amazonQProfileRegion(QRegionProfileManager.getInstance().activeProfile(project)?.region ?: "not-set")
40-
.ssoRegion(conn?.region)
41-
.credentialStartUrl(conn?.startUrl)
42-
.result(MetricResult.Failed)
43-
.reason(e.message)
44-
}
45-
throw e
46-
}
47-
?: error("Attempted to fetch profiles while there does not exist")
48-
val selectedProfile = QRegionProfileManager.getInstance().activeProfile(project) ?: profiles[0]
49-
ApplicationManager.getApplication().invokeLater {
50-
QRegionProfileDialog(
51-
project,
52-
profiles = profiles,
53-
selectedProfile = selectedProfile
54-
).show()
55-
}
56-
}
25+
QRegionProfileDialog(
26+
project,
27+
selectedProfile = QRegionProfileManager.getInstance().activeProfile(project)
28+
).show()
5729
}
5830
}

0 commit comments

Comments
 (0)