Skip to content

fix(amazonq): reauth workflow should not prompt profile selection page #5549

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 22 commits into from
Apr 10, 2025
Merged

fix(amazonq): reauth workflow should not prompt profile selection page #5549

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
8223470
Merge staging into feature/q-region-expansion
aws-toolkit-automation Apr 9, 2025
9105243
call if
evanliu048 Apr 9, 2025
ae016bd
add cache
evanliu048 Apr 9, 2025
c1474fe
linter
evanliu048 Apr 9, 2025
e2967bf
Merge staging into feature/q-region-expansion
aws-toolkit-automation Apr 9, 2025
cbb548d
Merge branch 'main' into regionExpansion_fixlaunch
evanliu048 Apr 9, 2025
0ec0d1f
Merge staging into feature/q-region-expansion
aws-toolkit-automation Apr 9, 2025
07e63ca
delete ut
evanliu048 Apr 9, 2025
de7a0f0
test
evanliu048 Apr 9, 2025
1a4a5db
Merge branch 'feature/q-region-expansion' into regionExpansion_fixlaunch
evanliu048 Apr 9, 2025
daff262
Merge branch 'main' into regionExpansion_fixlaunch
evanliu048 Apr 9, 2025
24f57fb
Merge branch 'main' into regionExpansion_fixlaunch
evanliu048 Apr 10, 2025
7a35aae
revise auth
evanliu048 Apr 10, 2025
b5d3299
re
evanliu048 Apr 10, 2025
29b297a
delete
evanliu048 Apr 10, 2025
3c4d39a
revert
evanliu048 Apr 10, 2025
62ddb32
Merge branch 'main' into regionExpansion_fixlaunch
evanliu048 Apr 10, 2025
bbc7cce
ut
evanliu048 Apr 10, 2025
ad24bbd
merge
evanliu048 Apr 10, 2025
fbcdf40
Merge branch 'main' into regionExpansion_fixlaunch
evanliu048 Apr 10, 2025
d3c76bd
reintialize q ui
evanliu048 Apr 10, 2025
5f97d38
Merge branch 'main' into regionExpansion_fixlaunch
rli Apr 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion ...c/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindowFactory.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,8 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware {
object : BearerTokenProviderListener {
override fun onChange(providerId: String, newScopes: List<String>?) {
if (ToolkitConnectionManager.getInstance(project).connectionStateForFeature(QConnection.getInstance()) == BearerTokenAuthState.AUTHORIZED) {
prepareChatContent(project, qPanel)
AmazonQToolWindow.getInstance(project).disposeAndRecreate()
qPanel.setContent(AmazonQToolWindow.getInstance(project).component)
}
}
}
Expand Down
34 changes: 33 additions & 1 deletion ...y/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ package software.aws.toolkits.jetbrains.services.codewhisperer
import com.intellij.openapi.project.Project
import com.intellij.testFramework.DisposableRule
import com.intellij.testFramework.ProjectRule
import com.intellij.testFramework.replaceService
import com.intellij.util.xmlb.XmlSerializer
import org.assertj.core.api.Assertions.assertThat
import org.jdom.output.XMLOutputter
Expand All @@ -15,7 +16,9 @@ import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.doReturn
import org.mockito.kotlin.mock
import org.mockito.kotlin.spy
import org.mockito.kotlin.stub
import org.mockito.kotlin.whenever
import software.amazon.awssdk.core.pagination.sync.SdkIterable
import software.amazon.awssdk.regions.Region
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
Expand All @@ -34,6 +37,7 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.logoutFromSsoConnection
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenAuthState
import software.aws.toolkits.jetbrains.core.region.MockRegionProviderRule
import software.aws.toolkits.jetbrains.services.amazonq.profile.QEndpoints
import software.aws.toolkits.jetbrains.services.amazonq.profile.QProfileResources
Expand Down Expand Up @@ -85,6 +89,10 @@ class QRegionProfileManagerTest {
sut = QRegionProfileManager()
val conn = authRule.createConnection(ManagedSsoProfile(ssoRegion = "us-east-1", startUrl = "", scopes = Q_SCOPES))
ToolkitConnectionManager.getInstance(project).switchConnection(conn)
val realManager = ToolkitConnectionManager.getInstance(project)
val managerSpy = spy(realManager)
doReturn(BearerTokenAuthState.AUTHORIZED).whenever(managerSpy).connectionStateForFeature(QConnection.getInstance())
project.replaceService(ToolkitConnectionManager::class.java, managerSpy, disposableRule.disposable)
}

@Test
Expand All @@ -106,7 +114,7 @@ class QRegionProfileManagerTest {
logoutFromSsoConnection(project, it)
}
}

ToolkitConnectionManager.getInstance(project).switchConnection(null)
assertThat(ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())).isNull()
assertThat(sut.activeProfile(project)).isNull()
}
Expand Down Expand Up @@ -384,4 +392,28 @@ class QRegionProfileManagerTest {
val profileList = actualState.connectionIdToProfileList["conn-123"]
assertThat(profileList).isEqualTo(2)
}

@Test
fun `getIdcConnectionOrNull handles NOT_AUTH and AUTHORIZED correctly`() {
val managerSpy = ToolkitConnectionManager.getInstance(project)
doReturn(BearerTokenAuthState.NOT_AUTHENTICATED).whenever(managerSpy)
.connectionStateForFeature(QConnection.getInstance())

// NOT AUTHORIZED
val notAuthConn = sut.getIdcConnectionOrNull(project)
assertThat(notAuthConn).isNull()

doReturn(BearerTokenAuthState.AUTHORIZED)
.whenever(managerSpy).connectionStateForFeature(QConnection.getInstance())

// AUTHORIZED
val normalConn = authRule.createConnection(
ManagedSsoProfile(ssoRegion = "us-east-1", startUrl = "", scopes = Q_SCOPES)
)
managerSpy.switchConnection(normalConn)

val normalConnectionResult = sut.getIdcConnectionOrNull(project)
assertThat(normalConnectionResult).isNotNull()
assertThat(normalConnectionResult).isEqualTo(normalConn)
}
}
15 changes: 10 additions & 5 deletions ...ity/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ 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.isSono
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenAuthState
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
import software.aws.toolkits.jetbrains.utils.notifyInfo
Expand Down Expand Up @@ -190,12 +191,16 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
return client
}

private fun getIdcConnectionOrNull(project: Project): AwsBearerTokenConnection? {
val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
if (connection is AwsBearerTokenConnection && !connection.isSono()) {
return connection
fun getIdcConnectionOrNull(project: Project): AwsBearerTokenConnection? {
val manager = ToolkitConnectionManager.getInstance(project)
val connection = manager.activeConnectionForFeature(QConnection.getInstance()) as? AwsBearerTokenConnection
Copy link
Contributor

@rli rli Apr 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i think at this point we need to completely remove CodeWhispererConnection if we are assuming everything is now QConnection

Copy link
Contributor Author

@evanliu048 evanliu048 Apr 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I may be missing part of the context, but I believe I'm using QConnection here—let me know if I'm misunderstanding anything

Copy link
Contributor

@Will-ShaoHua Will-ShaoHua Apr 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

because there was a old CodeWhispererConnection.kt before QConnection came in, so ideally we will check both connection for backward compatiblity. You can check isQConnected. I am in favor of removing codewhispererConnection as it would be a burden to check 2 connection type everywhere and considering we've been using Q_Scope for roughly 1 year, it shouldn't harm to remove CwsprConnection. I can follow up on the cwspr connectin deprecation if @rli you don't have concern about the deprecation

Copy link
Contributor Author

@evanliu048 evanliu048 Apr 10, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the clarification Will. We would need a separate pr to handle the codewhisperer connection deprecation right

Copy link
Contributor

@Will-ShaoHua Will-ShaoHua Apr 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes

val state = manager.connectionStateForFeature(QConnection.getInstance())

return if (connection != null && !connection.isSono() && state == BearerTokenAuthState.AUTHORIZED) {
connection
} else {
null
}
return null
}

companion object {
Expand Down
Loading