Skip to content

Commit 1a57043

Browse files
authored
fix(amazonq): reauth workflow should not prompt profile selection page (#5549)
### Problem On IDE startup, `QRegionProfileManager` loads the active profile from `aws.xml`. However, `validateProfile()` is called before SSO reauth completes, causing `listRegionProfiles()` to return empty. This results in the active profile being reset to `null`. ### Fix Update `getIdcConnectionOrNull()` to return the connection only if it's a valid Q connection **and the token is authorized**. Also, after successful reauth, immediately show the Q panel to avoid triggering the profile selection flow. ### Summary - Skip profile validation when token isn't ready - Ensure Q panel appears after reauth ### Changes - Refined `getIdcConnectionOrNull()` to check for `BearerTokenAuthState.AUTHORIZED` This should improve the IDE startup flow for users who rely on SSO, reducing friction and avoiding false-positive "profile selection required" states.
1 parent 330f565 commit 1a57043

File tree

3 files changed

+45
-7
lines changed

3 files changed

+45
-7
lines changed

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/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
2727
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
2828
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
2929
import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
30+
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenAuthState
3031
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
3132
import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
3233
import software.aws.toolkits.jetbrains.utils.notifyInfo
@@ -190,12 +191,16 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
190191
return client
191192
}
192193

193-
private fun getIdcConnectionOrNull(project: Project): AwsBearerTokenConnection? {
194-
val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
195-
if (connection is AwsBearerTokenConnection && !connection.isSono()) {
196-
return connection
194+
fun getIdcConnectionOrNull(project: Project): AwsBearerTokenConnection? {
195+
val manager = ToolkitConnectionManager.getInstance(project)
196+
val connection = manager.activeConnectionForFeature(QConnection.getInstance()) as? AwsBearerTokenConnection
197+
val state = manager.connectionStateForFeature(QConnection.getInstance())
198+
199+
return if (connection != null && !connection.isSono() && state == BearerTokenAuthState.AUTHORIZED) {
200+
connection
201+
} else {
202+
null
197203
}
198-
return null
199204
}
200205

201206
companion object {

0 commit comments

Comments
 (0)