diff --git a/.changes/3.63.json b/.changes/3.63.json new file mode 100644 index 00000000000..806e763d55a --- /dev/null +++ b/.changes/3.63.json @@ -0,0 +1,11 @@ +{ + "date" : "2025-04-08", + "version" : "3.63", + "entries" : [ { + "type" : "feature", + "description" : "Enterprise users can choose their preferred Amazon Q profile to improve personalization and workflow across different business regions" + }, { + "type" : "bugfix", + "description" : "Amazon Q /doc: close diff tab and open README file in preview mode after user accept changes" + } ] +} \ No newline at end of file diff --git a/.changes/3.64.json b/.changes/3.64.json new file mode 100644 index 00000000000..942136a9fd8 --- /dev/null +++ b/.changes/3.64.json @@ -0,0 +1,8 @@ +{ + "date" : "2025-04-10", + "version" : "3.64", + "entries" : [ { + "type" : "bugfix", + "description" : "Fix issue where IDE freezes when logging into Amazon Q" + } ] +} \ No newline at end of file diff --git a/.changes/next-release/bugfix-00947041-108c-4ef1-bec6-eb749dae7c2f.json b/.changes/next-release/bugfix-00947041-108c-4ef1-bec6-eb749dae7c2f.json new file mode 100644 index 00000000000..1406bab454a --- /dev/null +++ b/.changes/next-release/bugfix-00947041-108c-4ef1-bec6-eb749dae7c2f.json @@ -0,0 +1,4 @@ +{ + "type" : "bugfix", + "description" : "Fix issue where Amazon Q cannot process chunks from local `@workspace` context" +} \ No newline at end of file diff --git a/.changes/next-release/bugfix-19118cf8-9378-4bd6-bf5e-7e57520181d0.json b/.changes/next-release/bugfix-19118cf8-9378-4bd6-bf5e-7e57520181d0.json deleted file mode 100644 index 4121214a68a..00000000000 --- a/.changes/next-release/bugfix-19118cf8-9378-4bd6-bf5e-7e57520181d0.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "type" : "bugfix", - "description" : "Amazon Q /doc: close diff tab and open README file in preview mode after user accept changes" -} \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 321e017b9fe..b9932adcdf7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,10 @@ +# _3.64_ (2025-04-10) +- **(Bug Fix)** Fix issue where IDE freezes when logging into Amazon Q + +# _3.63_ (2025-04-08) +- **(Feature)** Enterprise users can choose their preferred Amazon Q profile to improve personalization and workflow across different business regions +- **(Bug Fix)** Amazon Q /doc: close diff tab and open README file in preview mode after user accept changes + # _3.62_ (2025-04-03) - **(Feature)** /review: automatically generate fix without clicking Generate Fix button - **(Bug Fix)** /transform: prompt user to re-authenticate if credentials expire during transformation diff --git a/gradle.properties b/gradle.properties index 5921806d5e7..616fef29d9b 100644 --- a/gradle.properties +++ b/gradle.properties @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # Toolkit Version -toolkitVersion=3.63-SNAPSHOT +toolkitVersion=3.65-SNAPSHOT # Publish Settings publishToken= diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/clients/AmazonQCodeGenerateClient.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/clients/AmazonQCodeGenerateClient.kt index 2c28c3240f3..3957d9fdf46 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/clients/AmazonQCodeGenerateClient.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/common/clients/AmazonQCodeGenerateClient.kt @@ -31,10 +31,10 @@ import software.aws.toolkits.core.utils.error import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.info import software.aws.toolkits.jetbrains.common.session.Intent -import software.aws.toolkits.jetbrains.core.awsClient import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.amazonqDoc.FEATURE_EVALUATION_PRODUCT_NAME import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTotalLatency import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata @@ -72,7 +72,7 @@ class AmazonQCodeGenerateClient(private val project: Project) { fun connection() = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) ?: error("Attempted to use connection while one does not exist") - fun bearerClient() = connection().getConnectionSettings().awsClient() + fun bearerClient() = QRegionProfileManager.getInstance().getQClient(project) private val amazonQStreamingClient get() = AmazonQStreamingClient.getInstance(project) @@ -88,6 +88,7 @@ class AmazonQCodeGenerateClient(private val project: Project) { } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(docUserContext) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } fun sendDocMetricData(operationName: String, result: String): SendTelemetryEventResponse = @@ -118,7 +119,9 @@ class AmazonQCodeGenerateClient(private val project: Project) { } fun createTaskAssistConversation(): CreateTaskAssistConversationResponse = bearerClient().createTaskAssistConversation( - CreateTaskAssistConversationRequest.builder().build() + CreateTaskAssistConversationRequest.builder() + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) + .build() ) fun createTaskAssistUploadUrl(conversationId: String, contentChecksumSha256: String, contentLength: Long): CreateUploadUrlResponse = @@ -137,6 +140,7 @@ class AmazonQCodeGenerateClient(private val project: Project) { ) .build() ) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } fun startTaskAssistCodeGeneration(conversationId: String, uploadId: String, userMessage: String, intent: Intent): StartTaskAssistCodeGenerationResponse = @@ -155,6 +159,7 @@ class AmazonQCodeGenerateClient(private val project: Project) { .uploadId(uploadId) } .intent(intent.name) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } fun getTaskAssistCodeGeneration(conversationId: String, codeGenerationId: String): GetTaskAssistCodeGenerationResponse = bearerClient() diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt index bc594209fba..3d3ea0aa3e3 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt @@ -6,6 +6,7 @@ package software.aws.toolkits.jetbrains.services.amazonq import com.intellij.openapi.Disposable import com.intellij.openapi.actionSystem.AnActionEvent import com.intellij.openapi.actionSystem.DataContext +import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.application.runInEdt import com.intellij.openapi.components.Service import com.intellij.openapi.components.service @@ -17,6 +18,7 @@ import com.intellij.ui.dsl.builder.Align import com.intellij.ui.dsl.builder.panel import com.intellij.ui.jcef.JBCefJSQuery import org.cef.CefApp +import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.error import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.warn @@ -33,11 +35,16 @@ import software.aws.toolkits.jetbrains.core.webview.BrowserState import software.aws.toolkits.jetbrains.core.webview.LoginBrowser import software.aws.toolkits.jetbrains.core.webview.WebviewResourceHandlerFactory import software.aws.toolkits.jetbrains.isDeveloperMode +import software.aws.toolkits.jetbrains.services.amazonq.profile.QProfileSwitchIntent +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.amazonq.util.createBrowser import software.aws.toolkits.jetbrains.utils.isQConnected import software.aws.toolkits.jetbrains.utils.isQExpired import software.aws.toolkits.jetbrains.utils.isQWebviewsAvailable import software.aws.toolkits.telemetry.FeatureId +import software.aws.toolkits.telemetry.MetricResult +import software.aws.toolkits.telemetry.Telemetry import software.aws.toolkits.telemetry.UiTelemetry import software.aws.toolkits.telemetry.WebviewTelemetry import java.awt.event.ActionListener @@ -204,6 +211,18 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos UiTelemetry.click(project, signInOption) } } + + is BrowserMessage.SwitchProfile -> { + QRegionProfileManager.getInstance().switchProfile( + project, + QRegionProfile(profileName = message.profileName, arn = message.arn), + intent = QProfileSwitchIntent.Auth + ) + } + + is BrowserMessage.PublishWebviewTelemetry -> { + publishTelemetry(message) + } } } @@ -244,13 +263,42 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos } // TODO: pass "REAUTH" if connection expires - val stage = if (isQExpired(project)) { - "REAUTH" - } else { - "START" - } + // Perform the potentially blocking AWS call outside the EDT to fetch available region profiles. + ApplicationManager.getApplication().executeOnPooledThread { + val stage = if (isQExpired(project)) { + "REAUTH" + } else if (isQConnected(project) && QRegionProfileManager.getInstance().isPendingProfileSelection(project)) { + "PROFILE_SELECT" + } else { + "START" + } - val jsonData = """ + var errorMessage: String? = null + var profiles: List = emptyList() + + if (stage == "PROFILE_SELECT") { + try { + profiles = QRegionProfileManager.getInstance().listRegionProfiles(project).orEmpty() + if (profiles.size == 1) { + LOG.debug { "User only have access to 1 Q profile, auto-selecting profile ${profiles.first().profileName} for ${project.name}" } + QRegionProfileManager.getInstance().switchProfile(project, profiles.first(), QProfileSwitchIntent.Update) + } + } catch (e: Exception) { + errorMessage = e.message + LOG.warn { "Failed to call listRegionProfiles API" } + val qConn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) + Telemetry.amazonq.didSelectProfile.use { span -> + span.source(QProfileSwitchIntent.Auth.value) + .amazonQProfileRegion(QRegionProfileManager.getInstance().activeProfile(project)?.region ?: "not-set") + .ssoRegion((qConn as? AwsBearerTokenConnection)?.region) + .credentialStartUrl((qConn as? AwsBearerTokenConnection)?.startUrl) + .result(MetricResult.Failed) + .reason(e.message) + } + } + } + + val jsonData = """ { stage: '$stage', regions: $regions, @@ -261,10 +309,16 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos }, cancellable: ${state.browserCancellable}, feature: '${state.feature}', - existConnections: ${writeValueAsString(selectionSettings.values.map { it.currentSelection }.toList())} + existConnections: ${writeValueAsString(selectionSettings.values.map { it.currentSelection }.toList())}, + profiles: ${writeValueAsString(profiles)}, + errorMessage: ${errorMessage?.let { "\"$it\"" } ?: "null"} + } + """.trimIndent() + + runInEdt { + executeJS("window.ideClient.prepareUi($jsonData)") } - """.trimIndent() - executeJS("window.ideClient.prepareUi($jsonData)") + } } override fun loginIAM(profileName: String, accessKey: String, secretKey: String) { diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/startup/AmazonQStartupActivity.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/startup/AmazonQStartupActivity.kt index b6b9aa29c9d..b85d94db10f 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/startup/AmazonQStartupActivity.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/startup/AmazonQStartupActivity.kt @@ -20,6 +20,7 @@ import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection import software.aws.toolkits.jetbrains.core.gettingstarted.emitUserState import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextController import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindow import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindowFactory @@ -53,6 +54,9 @@ class AmazonQStartupActivity : ProjectActivity { CodeWhispererExplorerActionManager.getInstance().setIsFirstRestartAfterQInstall(false) } } + + QRegionProfileManager.getInstance().validateProfile(project) + AmazonQLspService.getInstance(project) startLsp(project) if (runOnce.get()) return diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindow.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindow.kt index 15a533d1191..8a1d6378565 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindow.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindow.kt @@ -24,6 +24,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage import software.aws.toolkits.jetbrains.services.amazonq.messages.MessageConnector import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteraction import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteractionType +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.amazonq.util.highlightCommand import software.aws.toolkits.jetbrains.services.amazonq.webview.BrowserConnector import software.aws.toolkits.jetbrains.services.amazonq.webview.FqnWebviewAdapter @@ -127,7 +128,8 @@ class AmazonQToolWindow private constructor( isCodeScanAvailable = isCodeScanAvailable(project), isCodeTestAvailable = isCodeTestAvailable(project), isDocAvailable = isDocAvailable(project), - highlightCommand = highlightCommand() + highlightCommand = highlightCommand(), + activeProfile = QRegionProfileManager.getInstance().takeIf { it.shouldDisplayProfileInfo(project) }?.activeProfile(project) ) scope.launch { diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindowFactory.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindowFactory.kt index 3435075caaf..7137d4f966c 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindowFactory.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindowFactory.kt @@ -14,12 +14,10 @@ import com.intellij.ui.components.panels.Wrapper import com.intellij.util.ui.components.BorderLayoutPanel import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger -import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener 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.credentials.sso.bearer.BearerTokenProviderListener import software.aws.toolkits.jetbrains.core.notifications.NotificationPanel @@ -28,6 +26,9 @@ import software.aws.toolkits.jetbrains.core.webview.BrowserState import software.aws.toolkits.jetbrains.services.amazonq.QWebviewPanel import software.aws.toolkits.jetbrains.services.amazonq.RefreshQChatPanelButtonPressedListener import software.aws.toolkits.jetbrains.services.amazonq.gettingstarted.openMeetQPage +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener import software.aws.toolkits.jetbrains.utils.isQConnected import software.aws.toolkits.jetbrains.utils.isQExpired import software.aws.toolkits.jetbrains.utils.isQWebviewsAvailable @@ -62,7 +63,10 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware { ToolkitConnectionManagerListener.TOPIC, object : ToolkitConnectionManagerListener { override fun activeConnectionChanged(newConnection: ToolkitConnection?) { - onConnectionChanged(project, newConnection, qPanel) + ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let { qConn -> + openMeetQPage(project) + } + prepareChatContent(project, qPanel) } } ) @@ -71,9 +75,7 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware { RefreshQChatPanelButtonPressedListener.TOPIC, object : RefreshQChatPanelButtonPressedListener { override fun onRefresh() { - runInEdt { - prepareChatContent(project, qPanel) - } + prepareChatContent(project, qPanel) } } ) @@ -83,16 +85,23 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware { object : BearerTokenProviderListener { override fun onChange(providerId: String, newScopes: List?) { if (ToolkitConnectionManager.getInstance(project).connectionStateForFeature(QConnection.getInstance()) == BearerTokenAuthState.AUTHORIZED) { - val qComponent = AmazonQToolWindow.getInstance(project).component - - runInEdt { - qPanel.setContent(qComponent) - } + prepareChatContent(project, qPanel) } } } ) + project.messageBus.connect(toolWindow.disposable).subscribe( + QRegionProfileSelectedListener.TOPIC, + object : QRegionProfileSelectedListener { + override fun onProfileSelected(project: Project, profile: QRegionProfile?) { + if (project.isDisposed) return + AmazonQToolWindow.getInstance(project).disposeAndRecreate() + qPanel.setContent(AmazonQToolWindow.getInstance(project).component) + } + } + ) + prepareChatContent(project, qPanel) val content = contentManager.factory.createContent(mainPanel, null, false).also { @@ -107,13 +116,21 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware { project: Project, qPanel: Wrapper, ) { - val component = if (isQConnected(project) && !isQExpired(project)) { + /** + * only render Q Chat when + * 1. There is a Q connection + * 2. Q connection is not expired + * 3. User is not pending region profile selection + */ + val component = if (isQConnected(project) && !isQExpired(project) && !QRegionProfileManager.getInstance().isPendingProfileSelection(project)) { AmazonQToolWindow.getInstance(project).component } else { QWebviewPanel.getInstance(project).browser?.prepareBrowser(BrowserState(FeatureId.AmazonQ)) QWebviewPanel.getInstance(project).component } - qPanel.setContent(component) + runInEdt { + qPanel.setContent(component) + } } override fun init(toolWindow: ToolWindow) { @@ -134,36 +151,6 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware { override fun shouldBeAvailable(project: Project): Boolean = isQWebviewsAvailable() - private fun onConnectionChanged(project: Project, newConnection: ToolkitConnection?, qPanel: Wrapper) { - val isNewConnectionForQ = newConnection?.let { - (it as? AwsBearerTokenConnection)?.let { conn -> - val scopeShouldHave = Q_SCOPES - - LOG.debug { "newConnection: ${conn.id}; scope: ${conn.scopes}; scope must-have: $scopeShouldHave" } - - scopeShouldHave.all { s -> s in conn.scopes } - } ?: false - } ?: false - - if (isNewConnectionForQ) { - openMeetQPage(project) - } - - QWebviewPanel.getInstance(project).browser?.prepareBrowser(BrowserState(FeatureId.AmazonQ)) - - // isQConnected alone is not robust and there is race condition (read/update connection states) - val component = if (isNewConnectionForQ || (isQConnected(project) && !isQExpired(project))) { - LOG.debug { "returning Q-chat window; isQConnection=$isNewConnectionForQ; hasPinnedConnection=$isNewConnectionForQ" } - AmazonQToolWindow.getInstance(project).component - } else { - LOG.debug { "returning login window; no Q connection found" } - QWebviewPanel.getInstance(project).component - } - runInEdt { - qPanel.setContent(component) - } - } - companion object { private val LOG = getLogger() const val WINDOW_ID = AMAZON_Q_WINDOW_ID diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt index cc39985dd36..7dd41c795f1 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt @@ -8,6 +8,7 @@ import com.intellij.openapi.Disposable import com.intellij.openapi.util.Disposer import com.intellij.ui.jcef.JBCefJSQuery import org.cef.CefApp +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile import software.aws.toolkits.jetbrains.services.amazonq.util.HighlightCommand import software.aws.toolkits.jetbrains.services.amazonq.util.createBrowser import software.aws.toolkits.jetbrains.settings.MeetQSettings @@ -28,6 +29,7 @@ class Browser(parent: Disposable) : Disposable { isCodeScanAvailable: Boolean, isCodeTestAvailable: Boolean, highlightCommand: HighlightCommand?, + activeProfile: QRegionProfile?, ) { // register the scheme handler to route http://mynah/ URIs to the resources/assets directory on classpath CefApp.getInstance() @@ -37,7 +39,7 @@ class Browser(parent: Disposable) : Disposable { AssetResourceHandler.AssetResourceHandlerFactory(), ) - loadWebView(isCodeTransformAvailable, isFeatureDevAvailable, isDocAvailable, isCodeScanAvailable, isCodeTestAvailable, highlightCommand) + loadWebView(isCodeTransformAvailable, isFeatureDevAvailable, isDocAvailable, isCodeScanAvailable, isCodeTestAvailable, highlightCommand, activeProfile) } override fun dispose() { @@ -59,13 +61,22 @@ class Browser(parent: Disposable) : Disposable { isCodeScanAvailable: Boolean, isCodeTestAvailable: Boolean, highlightCommand: HighlightCommand?, + activeProfile: QRegionProfile?, ) { // setup empty state. The message request handlers use this for storing state // that's persistent between page loads. jcefBrowser.setProperty("state", "") // load the web app jcefBrowser.loadHTML( - getWebviewHTML(isCodeTransformAvailable, isFeatureDevAvailable, isDocAvailable, isCodeScanAvailable, isCodeTestAvailable, highlightCommand) + getWebviewHTML( + isCodeTransformAvailable, + isFeatureDevAvailable, + isDocAvailable, + isCodeScanAvailable, + isCodeTestAvailable, + highlightCommand, + activeProfile + ) ) } @@ -80,6 +91,7 @@ class Browser(parent: Disposable) : Disposable { isCodeScanAvailable: Boolean, isCodeTestAvailable: Boolean, highlightCommand: HighlightCommand?, + activeProfile: QRegionProfile?, ): String { val postMessageToJavaJsCode = receiveMessageQuery.inject("JSON.stringify(message)") @@ -100,7 +112,8 @@ class Browser(parent: Disposable) : Disposable { $isDocAvailable, // whether /doc is available $isCodeScanAvailable, // whether /scan is available $isCodeTestAvailable, // whether /test is available - ${OBJECT_MAPPER.writeValueAsString(highlightCommand)} + ${OBJECT_MAPPER.writeValueAsString(highlightCommand)}, + "${activeProfile?.profileName.orEmpty()}" ); } diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeScan/CodeScanChatApp.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeScan/CodeScanChatApp.kt index 166223d623a..0aa8dc42b04 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeScan/CodeScanChatApp.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeScan/CodeScanChatApp.kt @@ -21,6 +21,8 @@ import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQApp import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener import software.aws.toolkits.jetbrains.services.amazonqCodeScan.auth.isCodeScanAvailable import software.aws.toolkits.jetbrains.services.amazonqCodeScan.commands.CodeScanActionMessage import software.aws.toolkits.jetbrains.services.amazonqCodeScan.commands.CodeScanMessageListener @@ -141,6 +143,15 @@ class CodeScanChatApp(private val scope: CoroutineScope) : AmazonQApp { } } ) + + context.project.messageBus.connect(this).subscribe( + QRegionProfileSelectedListener.TOPIC, + object : QRegionProfileSelectedListener { + override fun onProfileSelected(project: Project, profile: QRegionProfile?) { + chatSessionStorage.deleteAllSessions() + } + } + ) } private fun getQTokenProvider(project: Project) = ( diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeScan/storage/ChatSessionStorage.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeScan/storage/ChatSessionStorage.kt index fb05a7beda8..57dcb85f44a 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeScan/storage/ChatSessionStorage.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeScan/storage/ChatSessionStorage.kt @@ -30,4 +30,8 @@ class ChatSessionStorage { fun changeAuthenticationNeededNotified(authNeededNotified: Boolean) { sessions.keys.forEach { sessions[it]?.authNeededNotified = authNeededNotified } } + + fun deleteAllSessions() { + sessions.clear() + } } diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatApp.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatApp.kt index 9da12272cdc..7972e45fb9c 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatApp.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeTestChatApp.kt @@ -4,6 +4,7 @@ package software.aws.toolkits.jetbrains.services.amazonqCodeTest import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.project.Project import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.launch import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection @@ -11,6 +12,8 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQApp import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener import software.aws.toolkits.jetbrains.services.amazonqCodeScan.auth.isCodeScanAvailable import software.aws.toolkits.jetbrains.services.amazonqCodeTest.auth.isCodeTestAvailable import software.aws.toolkits.jetbrains.services.amazonqCodeTest.controller.CodeTestChatController @@ -71,6 +74,15 @@ class CodeTestChatApp(private val scope: CoroutineScope) : AmazonQApp { } } ) + + context.project.messageBus.connect(this).subscribe( + QRegionProfileSelectedListener.TOPIC, + object : QRegionProfileSelectedListener { + override fun onProfileSelected(project: Project, profile: QRegionProfile?) { + chatSessionStorage.deleteAllSessions() + } + } + ) } private suspend fun handleMessage(message: AmazonQMessage, inboundAppMessagesHandler: InboundAppMessagesHandler) { diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt index dc60a60e747..aecede14c15 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt @@ -541,8 +541,8 @@ Please see the unit tests generated below. Click 'View Diff' to review the chang jobGroup = session.testGenerationJobGroupName, jobId = session.testGenerationJob, result = if (e.message == message("testgen.message.cancelled")) MetricResult.Cancelled else MetricResult.Failed, - reason = (e as CodeTestException).code ?: "DefaultError", - reasonDesc = if (e.message == message("testgen.message.cancelled")) "${e.code}: ${e.message}" else e.message, + reason = (e as? CodeTestException)?.code ?: "DefaultError", + reasonDesc = if (e.message == message("testgen.message.cancelled")) "${(e as? CodeTestException)?.code}: ${e.message}" else e.message, perfClientLatency = (Instant.now().toEpochMilli() - session.startTimeOfTestGeneration), isCodeBlockSelected = session.isCodeBlockSelected, artifactsUploadDuration = session.artifactUploadDuration, diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt index 3604ebd3549..f2906970f20 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt @@ -50,13 +50,11 @@ import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.info import software.aws.toolkits.core.utils.warn -import software.aws.toolkits.jetbrains.core.AwsClientManager import software.aws.toolkits.jetbrains.core.coroutines.EDT -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.isInternalUser import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.amazonq.project.RelevantDocument import software.aws.toolkits.jetbrains.services.amazonqCodeTest.CodeWhispererUTGChatManager import software.aws.toolkits.jetbrains.services.amazonqCodeTest.ConversationState @@ -277,9 +275,6 @@ class CodeTestChatController( promptInputDisabledState = true, ) // Send Request to Sync UTG API - val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) - // this should never happen because it should have been handled upstream by [AuthController] - ?: error("connection was found to be null") val contextExtractor = ActiveFileContextExtractor.create(fqnWebviewAdapter = null, project = project) val activeFileContext = ActiveFileContext( fileContext = FileContext( @@ -301,7 +296,7 @@ class CodeTestChatController( useRelevantDocuments = false, ) - val client = AwsClientManager.getInstance().getClient(connection.getConnectionSettings()) + val client = QRegionProfileManager.getInstance().getQClient(project) val request = requestData.toChatRequest() client.generateAssistantResponse(request, responseHandler).await() // TODO: Need to send isCodeBlockSelected field diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/storage/ChatSessionStorage.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/storage/ChatSessionStorage.kt index 0e38f06c6de..295bd9abe15 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/storage/ChatSessionStorage.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/storage/ChatSessionStorage.kt @@ -17,4 +17,8 @@ class ChatSessionStorage { // Find all sessions that are currently waiting to be authenticated fun getAuthenticatingSessions(): List = this.sessions.values.filter { it.isAuthenticating } + + fun deleteAllSessions() { + sessions.clear() + } } diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/DocApp.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/DocApp.kt index 0e55a655579..ed913374e08 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/DocApp.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/DocApp.kt @@ -4,6 +4,7 @@ package software.aws.toolkits.jetbrains.services.amazonqDoc import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.project.Project import kotlinx.coroutines.launch import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection @@ -11,6 +12,8 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQApp import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener import software.aws.toolkits.jetbrains.services.amazonqCodeScan.auth.isCodeScanAvailable import software.aws.toolkits.jetbrains.services.amazonqCodeTest.auth.isCodeTestAvailable import software.aws.toolkits.jetbrains.services.amazonqDoc.auth.isDocAvailable @@ -74,6 +77,15 @@ class DocApp : AmazonQApp { } } ) + + context.project.messageBus.connect(this).subscribe( + QRegionProfileSelectedListener.TOPIC, + object : QRegionProfileSelectedListener { + override fun onProfileSelected(project: Project, profile: QRegionProfile?) { + chatSessionStorage.deleteAllSessions() + } + } + ) } private suspend fun handleMessage(message: AmazonQMessage, inboundAppMessagesHandler: InboundAppMessagesHandler) { diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/storage/ChatSessionStorage.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/storage/ChatSessionStorage.kt index 2344dac5c94..e54a6b1d7fb 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/storage/ChatSessionStorage.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqDoc/storage/ChatSessionStorage.kt @@ -23,4 +23,11 @@ class ChatSessionStorage { // Find all sessions that are currently waiting to be authenticated fun getAuthenticatingSessions(): List = this.sessions.values.filter { it.isAuthenticating } + + fun deleteAllSessions() { + sessions.values.forEach { session -> + session.sessionState.token?.cancel() + } + sessions.clear() + } } diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevApp.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevApp.kt index 7169e391507..c9fe08efc64 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevApp.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevApp.kt @@ -4,6 +4,7 @@ package software.aws.toolkits.jetbrains.services.amazonqFeatureDev import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.project.Project import kotlinx.coroutines.launch import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection @@ -11,6 +12,8 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQApp import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener import software.aws.toolkits.jetbrains.services.amazonqCodeScan.auth.isCodeScanAvailable import software.aws.toolkits.jetbrains.services.amazonqCodeTest.auth.isCodeTestAvailable import software.aws.toolkits.jetbrains.services.amazonqDoc.auth.isDocAvailable @@ -75,6 +78,15 @@ class FeatureDevApp : AmazonQApp { } } ) + + context.project.messageBus.connect(this).subscribe( + QRegionProfileSelectedListener.TOPIC, + object : QRegionProfileSelectedListener { + override fun onProfileSelected(project: Project, profile: QRegionProfile?) { + chatSessionStorage.deleteAllSessions() + } + } + ) } private suspend fun handleMessage(message: AmazonQMessage, inboundAppMessagesHandler: InboundAppMessagesHandler) { diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClient.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClient.kt index 2d73f5fb4ba..7b25fb3b00e 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClient.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/clients/FeatureDevClient.kt @@ -28,10 +28,8 @@ import software.amazon.awssdk.services.codewhispererstreaming.model.ExportIntent import software.aws.toolkits.core.utils.error import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.info -import software.aws.toolkits.jetbrains.core.awsClient -import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager -import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.FEATURE_EVALUATION_PRODUCT_NAME import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTotalLatency import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata @@ -70,11 +68,7 @@ class FeatureDevClient( .build() } - private fun connection() = - ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) - ?: error("Attempted to use connection while one does not exist") - - private fun bearerClient() = connection().getConnectionSettings().awsClient() + private fun bearerClient() = QRegionProfileManager.getInstance().getQClient(project) private val amazonQStreamingClient get() = AmazonQStreamingClient.getInstance(project) @@ -88,6 +82,7 @@ class FeatureDevClient( } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(featureDevUserContext) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } fun sendFeatureDevMetricData(operationName: String, result: String): SendTelemetryEventResponse = @@ -115,6 +110,7 @@ class FeatureDevClient( } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(featureDevUserContext) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } fun sendFeatureDevCodeGenerationEvent( @@ -133,6 +129,7 @@ class FeatureDevClient( } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(featureDevUserContext) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } fun sendFeatureDevCodeAcceptanceEvent( @@ -151,11 +148,14 @@ class FeatureDevClient( } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(featureDevUserContext) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } fun createTaskAssistConversation(): CreateTaskAssistConversationResponse = bearerClient().createTaskAssistConversation( - CreateTaskAssistConversationRequest.builder().build(), + CreateTaskAssistConversationRequest.builder() + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) + .build(), ) fun createTaskAssistUploadUrl( @@ -182,6 +182,7 @@ class FeatureDevClient( .build(), ).build(), ) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } fun startTaskAssistCodeGeneration( @@ -205,6 +206,7 @@ class FeatureDevClient( .uploadId(uploadId) }.codeGenerationId(codeGenerationId.toString()) .currentCodeGenerationId(currentCodeGenerationId) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } fun getTaskAssistCodeGeneration( @@ -216,6 +218,7 @@ class FeatureDevClient( it .conversationId(conversationId) .codeGenerationId(codeGenerationId) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } suspend fun exportTaskAssistResultArchive(conversationId: String): MutableList = diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt index 5481b28b76e..75abc80ebf1 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt @@ -32,6 +32,8 @@ import software.aws.toolkits.jetbrains.core.coroutines.EDT import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener import software.aws.toolkits.jetbrains.services.amazonq.project.RepoSizeError import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindowFactory import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CodeIterationLimitException @@ -94,6 +96,17 @@ class FeatureDevController( private val authController: AuthController = AuthController(), ) : InboundAppMessagesHandler { + init { + context.project.messageBus.connect().subscribe( + QRegionProfileSelectedListener.TOPIC, + object : QRegionProfileSelectedListener { + override fun onProfileSelected(project: Project, profile: QRegionProfile?) { + chatSessionStorage.deleteAllSessions() + } + } + ) + } + val messenger = context.messagesFromAppToUi val toolWindow = ToolWindowManager.getInstance(context.project).getToolWindow(AmazonQToolWindowFactory.WINDOW_ID) diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/storage/ChatSessionStorage.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/storage/ChatSessionStorage.kt index 216d9cd1c8b..ed0f7eab7be 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/storage/ChatSessionStorage.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/storage/ChatSessionStorage.kt @@ -23,4 +23,11 @@ class ChatSessionStorage { // Find all sessions that are currently waiting to be authenticated fun getAuthenticatingSessions(): List = this.sessions.values.filter { it.isAuthenticating } + + fun deleteAllSessions() { + sessions.values.forEach { session -> + session.sessionState.token?.cancel() + } + sessions.clear() + } } diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/App.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/App.kt index b206595f9e8..34bb6bca614 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/App.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/App.kt @@ -4,6 +4,7 @@ package software.aws.toolkits.jetbrains.services.cwc import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.project.Project import kotlinx.coroutines.flow.merge import kotlinx.coroutines.launch import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope @@ -12,6 +13,8 @@ import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQApp import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteraction +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener import software.aws.toolkits.jetbrains.services.amazonq.util.highlightCommand import software.aws.toolkits.jetbrains.services.cwc.commands.ActionRegistrar import software.aws.toolkits.jetbrains.services.cwc.commands.CodeScanIssueActionMessage @@ -75,6 +78,15 @@ class App : AmazonQApp { } } ) + + ApplicationManager.getApplication().messageBus.connect(this).subscribe( + QRegionProfileSelectedListener.TOPIC, + object : QRegionProfileSelectedListener { + override fun onProfileSelected(project: Project, profile: QRegionProfile?) { + inboundAppMessagesHandler.processSessionClear() + } + } + ) } private suspend fun handleMessage(message: AmazonQMessage, inboundAppMessagesHandler: InboundAppMessagesHandler) { diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/InboundAppMessagesHandler.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/InboundAppMessagesHandler.kt index ff8a12e70ec..bbcddd84552 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/InboundAppMessagesHandler.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/InboundAppMessagesHandler.kt @@ -34,4 +34,6 @@ interface InboundAppMessagesHandler { suspend fun processCodeScanIssueAction(message: CodeScanIssueActionMessage) suspend fun processLinkClick(message: IncomingCwcMessage.ClickedLink) + + fun processSessionClear() } diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/v1/ChatSessionV1.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/v1/ChatSessionV1.kt index a55d9cc2075..8b563e24b32 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/v1/ChatSessionV1.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/v1/ChatSessionV1.kt @@ -36,10 +36,8 @@ import software.amazon.awssdk.services.codewhispererstreaming.model.UserInputMes import software.amazon.awssdk.services.codewhispererstreaming.model.UserIntent import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.info -import software.aws.toolkits.jetbrains.core.AwsClientManager import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineBgContext -import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager -import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.amazonq.project.RelevantDocument import software.aws.toolkits.jetbrains.services.cwc.ChatConstants import software.aws.toolkits.jetbrains.services.cwc.clients.chat.ChatSession @@ -170,11 +168,7 @@ class ChatSessionV1( try { withTimeout(ChatConstants.REQUEST_TIMEOUT_MS.toLong()) { - val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) - // this should never happen because it should have been handled upstream by [AuthController] - ?: error("connection was found to be null") - - val client = AwsClientManager.getInstance().getClient(connection.getConnectionSettings()) + val client = QRegionProfileManager.getInstance().getQClient(project) val request = data.toChatRequest() logger.info { "Request from tab: ${data.tabId}, conversationId: $conversationId, request: $request" } client.generateAssistantResponse(request, responseHandler).await() diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/SendToQActionGroup.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/SendToQActionGroup.kt index d3e77515cd1..d903f842020 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/SendToQActionGroup.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/SendToQActionGroup.kt @@ -8,6 +8,7 @@ import com.intellij.openapi.actionSystem.AnActionEvent import com.intellij.openapi.actionSystem.DefaultActionGroup import com.intellij.openapi.project.DumbAware import com.intellij.openapi.wm.ToolWindowManager +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindowFactory class SendToQActionGroup : DefaultActionGroup(), DumbAware { @@ -16,6 +17,7 @@ class SendToQActionGroup : DefaultActionGroup(), DumbAware { override fun update(e: AnActionEvent) { val project = e.project ?: return val amazonQWindow = ToolWindowManager.getInstance(project).getToolWindow(AmazonQToolWindowFactory.WINDOW_ID) - e.presentation.isEnabledAndVisible = amazonQWindow?.isAvailable ?: false + e.presentation.isEnabledAndVisible = (amazonQWindow?.isAvailable == true) && + !QRegionProfileManager.getInstance().hasValidConnectionButNoActiveProfile(project) } } diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt index 25630d63bab..800d7f5dea4 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt @@ -541,6 +541,10 @@ class ChatController private constructor( .map { it.tabId } .first() + override fun processSessionClear() { + chatSessionStorage.deleteAllSessions() + } + companion object { private val logger = getLogger() diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/inline/InlineChatController.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/inline/InlineChatController.kt index 1a6f8d2d10f..bf468976a2f 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/inline/InlineChatController.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/inline/InlineChatController.kt @@ -55,6 +55,8 @@ import software.aws.toolkits.jetbrains.core.gettingstarted.requestCredentialsFor import software.aws.toolkits.jetbrains.core.webview.BrowserState import software.aws.toolkits.jetbrains.services.amazonq.QWebviewPanel import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AMAZON_Q_WINDOW_ID import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition @@ -101,6 +103,14 @@ class InlineChatController( init { Disposer.register(this, listener) project.messageBus.connect(this).subscribe(FileEditorManagerListener.FILE_EDITOR_MANAGER, listener) + project.messageBus.connect(this).subscribe( + QRegionProfileSelectedListener.TOPIC, + object : QRegionProfileSelectedListener { + override fun onProfileSelected(project: Project, profile: QRegionProfile?) { + sessionStorage.deleteAllSessions() + } + } + ) } data class InlineChatMetrics( diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/inline/OpenChatInputAction.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/inline/OpenChatInputAction.kt index ebc6b53cd13..f581a217226 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/inline/OpenChatInputAction.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/inline/OpenChatInputAction.kt @@ -3,10 +3,12 @@ package software.aws.toolkits.jetbrains.services.cwc.inline +import com.intellij.openapi.actionSystem.ActionUpdateThread import com.intellij.openapi.actionSystem.AnAction import com.intellij.openapi.actionSystem.AnActionEvent import com.intellij.openapi.actionSystem.CommonDataKeys import com.intellij.openapi.util.Key +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager class OpenChatInputAction : AnAction() { override fun actionPerformed(e: AnActionEvent) { @@ -24,4 +26,9 @@ class OpenChatInputAction : AnAction() { val inlineChatController = InlineChatController.getInstance(project) inlineChatController.initPopup(editor) } + override fun update(e: AnActionEvent) { + val project = e.project ?: return + e.presentation.isEnabledAndVisible = !QRegionProfileManager.getInstance().hasValidConnectionButNoActiveProfile(project) + } + override fun getActionUpdateThread() = ActionUpdateThread.BGT } diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/storage/ChatSessionStorage.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/storage/ChatSessionStorage.kt index 231cdd21d5f..e11d0693b7c 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/storage/ChatSessionStorage.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/storage/ChatSessionStorage.kt @@ -26,4 +26,8 @@ class ChatSessionStorage( fun deleteSession(tabId: String) { sessions.remove(tabId)?.scope?.cancel() } + + fun deleteAllSessions() { + sessions.clear() + } } diff --git a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerManager.kt b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerManager.kt index a827a4bab8d..9b2db462a41 100644 --- a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerManager.kt +++ b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerManager.kt @@ -32,6 +32,8 @@ import software.aws.toolkits.core.utils.warn import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope import software.aws.toolkits.jetbrains.services.amazonq.CODE_TRANSFORM_TROUBLESHOOT_DOC_MVN_FAILURE import software.aws.toolkits.jetbrains.services.amazonq.CODE_TRANSFORM_TROUBLESHOOT_DOC_PROJECT_SIZE +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener import software.aws.toolkits.jetbrains.services.codemodernizer.client.GumbyClient import software.aws.toolkits.jetbrains.services.codemodernizer.commands.CodeTransformMessageListener import software.aws.toolkits.jetbrains.services.codemodernizer.constants.HIL_POM_FILE_NAME @@ -130,6 +132,23 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo init { CodeModernizerSessionState.getInstance(project).setDefaults() + initQRegionProfileSelectedListener() + } + + private fun initQRegionProfileSelectedListener() { + project.messageBus.connect(this).subscribe( + QRegionProfileSelectedListener.TOPIC, + object : QRegionProfileSelectedListener { + override fun onProfileSelected(project: Project, profile: QRegionProfile?) { + stopModernize() + codeTransformationSession?.let { + Disposer.dispose(it) + } + managerState = CodeModernizerState() + codeTransformationSession = null + } + } + ) } fun validate(project: Project, transformationType: CodeTransformType): ValidationResult { @@ -259,7 +278,7 @@ class CodeModernizerManager(private val project: Project) : PersistentStateCompo fun runModernize(copyResult: MavenCopyCommandsResult? = null) { initStopParameters() - val session = codeTransformationSession as CodeModernizerSession + val session = codeTransformationSession ?: return initModernizationJobUI(true, project.getModuleOrProjectNameForFile(session.sessionContext.configurationFile)) launchModernizationJob(session, copyResult) } diff --git a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerSession.kt b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerSession.kt index 9ba17c22cb1..921a9ac0f2c 100644 --- a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerSession.kt +++ b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerSession.kt @@ -5,6 +5,7 @@ package software.aws.toolkits.jetbrains.services.codemodernizer import com.intellij.openapi.Disposable import com.intellij.openapi.application.runInEdt +import com.intellij.openapi.util.Disposer import com.intellij.serviceContainer.AlreadyDisposedException import com.intellij.util.io.HttpRequests import kotlinx.coroutines.delay @@ -149,6 +150,9 @@ class CodeModernizerSession( * Based on [CodeWhispererCodeScanSession] */ suspend fun createModernizationJob(copyResult: MavenCopyCommandsResult?): CodeModernizerStartJobResult { + if (this.isDisposed.get()) { + return CodeModernizerStartJobResult.Cancelled + } LOG.info { "Compressing local project" } val payload: File? var payloadSize = 0 @@ -182,6 +186,9 @@ class CodeModernizerSession( payloadSize = payload.length().toInt() LOG.info { "Uploading zip file with size: $payloadSize bytes" } + if (this.isDisposed.get()) { + return CodeModernizerStartJobResult.Cancelled + } if (payloadSize > MAX_ZIP_SIZE) { telemetryErrorMessage = "Project exceeds max upload size" @@ -210,7 +217,7 @@ class CodeModernizerSession( telemetryErrorMessage = "Credential expired before uploading project" return CodeModernizerStartJobResult.ZipUploadFailed(UploadFailureReason.CREDENTIALS_EXPIRED) } - if (shouldStop.get()) { + if (shouldStop.get() || this.isDisposed.get()) { LOG.warn { "Job was cancelled by user before upload was called" } telemetryErrorMessage = "Cancelled when about to upload project" return CodeModernizerStartJobResult.Cancelled @@ -285,7 +292,7 @@ class CodeModernizerSession( CodeTransformMessageListener.instance.onUploadResult() return try { - if (shouldStop.get()) { + if (shouldStop.get() || this.isDisposed.get()) { LOG.warn { "Job was cancelled by user before start job was called" } return CodeModernizerStartJobResult.Cancelled } @@ -625,6 +632,8 @@ class CodeModernizerSession( override fun dispose() { isDisposed.set(true) + shouldStop.set(true) + Disposer.dispose(sessionContext) } fun getActiveJobId() = state.currentJobId diff --git a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeTransformChatApp.kt b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeTransformChatApp.kt index 1c32f2f07c5..b102d348d11 100644 --- a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeTransformChatApp.kt +++ b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeTransformChatApp.kt @@ -4,6 +4,7 @@ package software.aws.toolkits.jetbrains.services.codemodernizer import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.project.Project import kotlinx.coroutines.flow.merge import kotlinx.coroutines.launch import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope @@ -15,6 +16,8 @@ import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQApp import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener import software.aws.toolkits.jetbrains.services.amazonqCodeScan.auth.isCodeScanAvailable import software.aws.toolkits.jetbrains.services.amazonqCodeTest.auth.isCodeTestAvailable import software.aws.toolkits.jetbrains.services.amazonqDoc.auth.isDocAvailable @@ -161,6 +164,15 @@ class CodeTransformChatApp : AmazonQApp { } } ) + + context.project.messageBus.connect(this).subscribe( + QRegionProfileSelectedListener.TOPIC, + object : QRegionProfileSelectedListener { + override fun onProfileSelected(project: Project, profile: QRegionProfile?) { + chatSessionStorage.deleteAllSessions() + } + } + ) } private suspend fun handleMessage(message: AmazonQMessage, inboundAppMessagesHandler: InboundAppMessagesHandler) { diff --git a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/client/GumbyClient.kt b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/client/GumbyClient.kt index 097b9231d4a..6134e69ca22 100644 --- a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/client/GumbyClient.kt +++ b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/client/GumbyClient.kt @@ -38,9 +38,6 @@ import software.aws.toolkits.core.utils.error import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.info import software.aws.toolkits.jetbrains.core.AwsClientManager -import software.aws.toolkits.jetbrains.core.awsClient -import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager -import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection import software.aws.toolkits.jetbrains.services.amazonq.APPLICATION_ZIP import software.aws.toolkits.jetbrains.services.amazonq.AWS_KMS import software.aws.toolkits.jetbrains.services.amazonq.CONTENT_SHA256 @@ -48,6 +45,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.SERVER_SIDE_ENCRYPTION import software.aws.toolkits.jetbrains.services.amazonq.SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient import software.aws.toolkits.jetbrains.services.amazonq.codeWhispererUserContext +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerMetrics import software.aws.toolkits.jetbrains.services.codemodernizer.model.JobId import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTotalLatency @@ -58,10 +56,7 @@ import java.time.Instant @Service(Service.Level.PROJECT) class GumbyClient(private val project: Project) { - private fun connection() = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) - ?: error("Attempted to use connection while one does not exist") - - private fun bearerClient() = connection().getConnectionSettings().awsClient() + private fun bearerClient() = QRegionProfileManager.getInstance().getQClient(project) private val amazonQStreamingClient get() = AmazonQStreamingClient.getInstance(project) @@ -71,6 +66,7 @@ class GumbyClient(private val project: Project) { .contentChecksumType(ContentChecksumType.SHA_256) .contentChecksum(sha256Checksum) .uploadIntent(UploadIntent.TRANSFORMATION) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) .build() return callApi({ bearerClient().createUploadUrl(request) }, apiName = "CreateUploadUrl") } @@ -92,12 +88,16 @@ class GumbyClient(private val project: Project) { ) .build() ) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) .build() return callApi({ bearerClient().createUploadUrl(request) }, apiName = "CreateUploadUrl") } fun getCodeModernizationJob(jobId: String): GetTransformationResponse { - val request = GetTransformationRequest.builder().transformationJobId(jobId).build() + val request = GetTransformationRequest.builder() + .transformationJobId(jobId) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) + .build() return callApi({ bearerClient().getTransformation(request) }, apiName = "GetTransformation") } @@ -116,6 +116,7 @@ class GumbyClient(private val project: Project) { .source { it.language(sourceLanguage) } .target { it.language(targetLanguage) } } + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) .build() return callApi({ bearerClient().startTransformation(request) }, apiName = "StartTransformation") } @@ -127,17 +128,22 @@ class GumbyClient(private val project: Project) { val request = ResumeTransformationRequest.builder() .transformationJobId(jobId.id) .userActionStatus(userActionStatus) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) .build() return callApi({ bearerClient().resumeTransformation(request) }, apiName = "ResumeTransformation") } fun getCodeModernizationPlan(jobId: JobId): GetTransformationPlanResponse { - val request = GetTransformationPlanRequest.builder().transformationJobId(jobId.id).build() + val request = GetTransformationPlanRequest.builder() + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) + .transformationJobId(jobId.id).build() return callApi({ bearerClient().getTransformationPlan(request) }, apiName = "GetTransformationPlan") } fun stopTransformation(transformationJobId: String): StopTransformationResponse { - val request = StopTransformationRequest.builder().transformationJobId(transformationJobId).build() + val request = StopTransformationRequest.builder().transformationJobId(transformationJobId) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) + .build() return callApi({ bearerClient().stopTransformation(request) }, apiName = "StopTransformation") } @@ -232,6 +238,7 @@ class GumbyClient(private val project: Project) { } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } } diff --git a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/ideMaven/MavenRunnerUtils.kt b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/ideMaven/MavenRunnerUtils.kt index 74ec6aeae3d..6629b8c2d34 100644 --- a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/ideMaven/MavenRunnerUtils.kt +++ b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/ideMaven/MavenRunnerUtils.kt @@ -16,6 +16,7 @@ import org.slf4j.Logger import software.aws.toolkits.core.utils.error import software.aws.toolkits.core.utils.info import software.aws.toolkits.jetbrains.services.codemodernizer.CodeTransformTelemetryManager +import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerSessionContext import software.aws.toolkits.jetbrains.services.codemodernizer.model.MavenCopyCommandsResult import software.aws.toolkits.jetbrains.services.codemodernizer.model.MavenDependencyReportCommandsResult import software.aws.toolkits.telemetry.CodeTransformBuildCommand @@ -25,6 +26,7 @@ import java.nio.file.Files import java.nio.file.Path fun runHilMavenCopyDependency( + context: CodeModernizerSessionContext, sourceFolder: File, destinationDir: File, logBuilder: StringBuilder, @@ -35,6 +37,7 @@ fun runHilMavenCopyDependency( try { // Create shared parameters val transformMvnRunner = TransformMavenRunner(project) + context.mavenRunnerQueue.add(transformMvnRunner) val mvnSettings = MavenRunner.getInstance(project).settings.clone() // clone required to avoid editing user settings // run copy dependencies @@ -57,7 +60,14 @@ fun runHilMavenCopyDependency( return MavenCopyCommandsResult.Success(destinationDir) } -fun runMavenCopyCommands(sourceFolder: File, logBuilder: StringBuilder, logger: Logger, project: Project, shouldSkipTests: Boolean): MavenCopyCommandsResult { +fun runMavenCopyCommands( + context: CodeModernizerSessionContext, + sourceFolder: File, + logBuilder: StringBuilder, + logger: Logger, + project: Project, + shouldSkipTests: Boolean, +): MavenCopyCommandsResult { val currentTimestamp = System.currentTimeMillis() val destinationDir = Files.createTempDirectory("transformation_dependencies_temp_$currentTimestamp") val telemetry = CodeTransformTelemetryManager.getInstance(project) @@ -68,6 +78,7 @@ fun runMavenCopyCommands(sourceFolder: File, logBuilder: StringBuilder, logger: try { // Create shared parameters val transformMvnRunner = TransformMavenRunner(project) + context.mavenRunnerQueue.add(transformMvnRunner) val mvnSettings = MavenRunner.getInstance(project).settings.clone() // clone required to avoid editing user settings val sourceVirtualFile = LocalFileSystem.getInstance().findFileByIoFile(sourceFolder) @@ -282,10 +293,17 @@ private fun runMavenDependencyUpdatesReport( return dependencyUpdatesReportRunnable } -fun runDependencyReportCommands(sourceFolder: File, logBuilder: StringBuilder, logger: Logger, project: Project): MavenDependencyReportCommandsResult { +fun runDependencyReportCommands( + context: CodeModernizerSessionContext, + sourceFolder: File, + logBuilder: StringBuilder, + logger: Logger, + project: Project, +): MavenDependencyReportCommandsResult { logger.info { "Executing IntelliJ bundled Maven" } val transformMvnRunner = TransformMavenRunner(project) + context.mavenRunnerQueue.add(transformMvnRunner) val mvnSettings = MavenRunner.getInstance(project).settings.clone() // clone required to avoid editing user settings val runnable = runMavenDependencyUpdatesReport(sourceFolder, logBuilder, mvnSettings, transformMvnRunner, logger) diff --git a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/ideMaven/TransformMavenRunner.kt b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/ideMaven/TransformMavenRunner.kt index f202cf6896f..a68c659ba6f 100644 --- a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/ideMaven/TransformMavenRunner.kt +++ b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/ideMaven/TransformMavenRunner.kt @@ -5,6 +5,7 @@ package software.aws.toolkits.jetbrains.services.codemodernizer.ideMaven import com.intellij.execution.process.ProcessAdapter import com.intellij.execution.process.ProcessEvent +import com.intellij.execution.process.ProcessHandler import com.intellij.execution.process.ProcessOutputTypes import com.intellij.execution.runners.ProgramRunner import com.intellij.execution.ui.RunContentDescriptor @@ -16,11 +17,13 @@ import org.jetbrains.idea.maven.execution.MavenRunnerParameters import org.jetbrains.idea.maven.execution.MavenRunnerSettings class TransformMavenRunner(val project: Project) { + private var handler: ProcessHandler? = null fun run(parameters: MavenRunnerParameters, settings: MavenRunnerSettings, onComplete: TransformRunnable) { FileDocumentManager.getInstance().saveAllDocuments() val callback = ProgramRunner.Callback { descriptor: RunContentDescriptor -> val handler = descriptor.processHandler + this.handler = handler if (handler == null) { // add log error here onComplete.setExitCode(-1) @@ -50,4 +53,8 @@ class TransformMavenRunner(val project: Project) { // Setting isDelegateBuild = true allows us to set the JRE used by Maven during runtime MavenRunConfigurationType.runConfiguration(project, parameters, null, settings, callback, false) } + + fun cancel() { + this.handler?.destroyProcess() + } } diff --git a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/model/CodeModernizerSessionContext.kt b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/model/CodeModernizerSessionContext.kt index 6a1a75b9d8a..a0da0824a76 100644 --- a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/model/CodeModernizerSessionContext.kt +++ b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/model/CodeModernizerSessionContext.kt @@ -4,6 +4,7 @@ package software.aws.toolkits.jetbrains.services.codemodernizer.model import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper +import com.intellij.openapi.Disposable import com.intellij.openapi.application.runInEdt import com.intellij.openapi.application.runReadAction import com.intellij.openapi.project.Project @@ -19,6 +20,7 @@ import software.aws.toolkits.core.utils.putNextEntry import software.aws.toolkits.jetbrains.services.codemodernizer.EXPLAINABILITY_V1 import software.aws.toolkits.jetbrains.services.codemodernizer.constants.HIL_DEPENDENCIES_ROOT_NAME import software.aws.toolkits.jetbrains.services.codemodernizer.constants.HIL_MANIFEST_FILE_NAME +import software.aws.toolkits.jetbrains.services.codemodernizer.ideMaven.TransformMavenRunner import software.aws.toolkits.jetbrains.services.codemodernizer.ideMaven.runDependencyReportCommands import software.aws.toolkits.jetbrains.services.codemodernizer.ideMaven.runHilMavenCopyDependency import software.aws.toolkits.jetbrains.services.codemodernizer.ideMaven.runMavenCopyCommands @@ -59,6 +61,7 @@ const val INVALID_SUFFIX_REPOSITORIES = "repositories" const val ORACLE_DB = "ORACLE" const val AURORA_DB = "AURORA_POSTGRESQL" const val RDS_DB = "POSTGRESQL" + data class CodeModernizerSessionContext( val project: Project, var configurationFile: VirtualFile? = null, // used to ZIP module @@ -71,9 +74,11 @@ data class CodeModernizerSessionContext( val sourceServerName: String? = null, var schema: String? = null, val sqlMetadataZip: File? = null, -) { +) : Disposable { private val mapper = jacksonObjectMapper() private val ignoredDependencyFileExtensions = setOf(INVALID_SUFFIX_SHA, INVALID_SUFFIX_REPOSITORIES) + private var isDisposed = false + val mavenRunnerQueue: MutableList = mutableListOf() private fun File.isMavenTargetFolder(): Boolean { val hasPomSibling = this.resolveSibling(MAVEN_CONFIGURATION_FILE_NAME).exists() @@ -100,19 +105,22 @@ data class CodeModernizerSessionContext( } fun executeMavenCopyCommands(sourceFolder: File, buildLogBuilder: StringBuilder): MavenCopyCommandsResult { + if (isDisposed) return MavenCopyCommandsResult.Cancelled val shouldSkipTests = customBuildCommand == MAVEN_BUILD_SKIP_UNIT_TESTS - return runMavenCopyCommands(sourceFolder, buildLogBuilder, LOG, project, shouldSkipTests) + return runMavenCopyCommands(this, sourceFolder, buildLogBuilder, LOG, project, shouldSkipTests) } private fun executeHilMavenCopyDependency(sourceFolder: File, destinationFolder: File, buildLogBuilder: StringBuilder) = runHilMavenCopyDependency( + this, sourceFolder, destinationFolder, buildLogBuilder, LOG, - project + project, ) fun copyHilDependencyUsingMaven(hilTepDirPath: Path): MavenCopyCommandsResult { + if (isDisposed) return MavenCopyCommandsResult.Cancelled val sourceFolder = File(getPathToHilArtifactPomFolder(hilTepDirPath).pathString) val destinationFolder = Files.createDirectories(getPathToHilDependenciesRootDir(hilTepDirPath)).toFile() val buildLogBuilder = StringBuilder("Starting Build Log...\n") @@ -121,6 +129,7 @@ data class CodeModernizerSessionContext( } fun getDependenciesUsingMaven(): MavenCopyCommandsResult { + if (isDisposed) return MavenCopyCommandsResult.Cancelled val root = configurationFile?.parent val sourceFolder = File(root?.path) val buildLogBuilder = StringBuilder("Starting Build Log...\n") @@ -128,14 +137,16 @@ data class CodeModernizerSessionContext( } fun createDependencyReportUsingMaven(hilTempPomPath: Path): MavenDependencyReportCommandsResult { + if (isDisposed) return MavenDependencyReportCommandsResult.Cancelled val sourceFolder = File(hilTempPomPath.pathString) val buildLogBuilder = StringBuilder("Starting Build Log...\n") return executeDependencyVersionReportUsingMaven(sourceFolder, buildLogBuilder) } + private fun executeDependencyVersionReportUsingMaven( sourceFolder: File, buildLogBuilder: StringBuilder, - ) = runDependencyReportCommands(sourceFolder, buildLogBuilder, LOG, project) + ) = runDependencyReportCommands(this, sourceFolder, buildLogBuilder, LOG, project) fun createZipForHilUpload(hilTempPath: Path, manifest: CodeTransformHilDownloadManifest?, targetVersion: String): ZipCreationResult = runReadAction { @@ -326,6 +337,13 @@ data class CodeModernizerSessionContext( CodeModernizerBottomWindowPanelManager.getInstance(project).setJobStartingUI() } + override fun dispose() { + isDisposed = true + this.mavenRunnerQueue.forEach { + it.cancel() + } + } + companion object { private val LOG = getLogger() } diff --git a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/session/ChatSessionStorage.kt b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/session/ChatSessionStorage.kt index 807e5dcbe1b..99198a56550 100644 --- a/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/session/ChatSessionStorage.kt +++ b/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/session/ChatSessionStorage.kt @@ -28,4 +28,8 @@ class ChatSessionStorage { fun changeAuthenticationNeededNotified(authNeededNotified: Boolean) { sessions.keys.forEach { sessions[it]?.authNeededNotified = authNeededNotified } } + + fun deleteAllSessions() { + sessions.clear() + } } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/resources/META-INF/plugin-codewhisperer.xml b/plugins/amazonq/codewhisperer/jetbrains-community/resources/META-INF/plugin-codewhisperer.xml index 07d8f29e713..f8a2de49534 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/resources/META-INF/plugin-codewhisperer.xml +++ b/plugins/amazonq/codewhisperer/jetbrains-community/resources/META-INF/plugin-codewhisperer.xml @@ -32,6 +32,7 @@ + + diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererRecommendationAction.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererRecommendationAction.kt index 97fcc101b8a..7eae90e6af0 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererRecommendationAction.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererRecommendationAction.kt @@ -11,6 +11,7 @@ import com.intellij.openapi.project.DumbAware import com.intellij.openapi.util.Key import kotlinx.coroutines.Job import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.codewhisperer.model.LatencyContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.TriggerTypeInfo import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType @@ -28,6 +29,10 @@ class CodeWhispererRecommendationAction : AnAction(message("codewhisperer.trigge } override fun actionPerformed(e: AnActionEvent) { + val project = e.project ?: return + if (QRegionProfileManager.getInstance().hasValidConnectionButNoActiveProfile(project)) { + return + } val latencyContext = LatencyContext() latencyContext.codewhispererPreprocessingStart = System.nanoTime() latencyContext.codewhispererEndToEndStart = System.nanoTime() diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/AmazonQCodeFixSession.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/AmazonQCodeFixSession.kt index 899fbd030a7..7ed08eb511f 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/AmazonQCodeFixSession.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/AmazonQCodeFixSession.kt @@ -25,6 +25,7 @@ import software.aws.toolkits.core.utils.error import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.info import software.aws.toolkits.core.utils.putNextEntry +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererZipUploadManager @@ -148,6 +149,7 @@ class AmazonQCodeFixSession(val project: Project) { .artifactType(artifactType) .uploadIntent(UploadIntent.CODE_FIX_GENERATION) .uploadContext(UploadContext.fromCodeFixUploadContext(CodeFixUploadContext.builder().codeFixName(codeFixName).build())) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) .build() ) } catch (e: Exception) { @@ -176,6 +178,7 @@ class AmazonQCodeFixSession(val project: Project) { .ruleId(ruleId) .description(description) .referenceTrackerConfiguration { it.recommendationsWithReferences(includeCodeWithReference) } + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) .build() return try { @@ -200,6 +203,7 @@ class AmazonQCodeFixSession(val project: Project) { val request = GetCodeFixJobRequest.builder() .jobId(jobId) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) .build() val response = clientAdaptor.getCodeFixJob(request) @@ -226,7 +230,11 @@ class AmazonQCodeFixSession(val project: Project) { } private fun getCodeFixJob(jobId: String): GetCodeFixJobResponse { - val response = clientAdaptor.getCodeFixJob(GetCodeFixJobRequest.builder().jobId(jobId).build()) + val response = clientAdaptor.getCodeFixJob( + GetCodeFixJobRequest.builder() + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) + .jobId(jobId).build() + ) return response } private fun zipFile(file: Path): File = createTemporaryZipFile { diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanSession.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanSession.kt index 0a7bb0aea62..27e50d5351b 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanSession.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanSession.kt @@ -33,6 +33,7 @@ import software.aws.toolkits.core.utils.Waiters.waitUntil import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.info +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.CodeScanSessionConfig import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.PayloadContext import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor @@ -272,6 +273,7 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) { .artifacts(artifactsMap) .scope(scope.value) .codeScanName(codeScanName) + .profileArn(QRegionProfileManager.getInstance().activeProfile(sessionContext.project)?.arn) .build() ) } catch (e: Exception) { @@ -285,6 +287,7 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) { clientAdaptor.getCodeScan( GetCodeAnalysisRequest.builder() .jobId(jobId) + .profileArn(QRegionProfileManager.getInstance().activeProfile(sessionContext.project)?.arn) .build() ) } catch (e: Exception) { @@ -299,6 +302,7 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) { .jobId(jobId) .codeAnalysisFindingsSchema(CodeAnalysisFindingsSchema.CODEANALYSIS_FINDINGS_1_0) .nextToken(nextToken) + .profileArn(QRegionProfileManager.getInstance().activeProfile(sessionContext.project)?.arn) .build() ) } catch (e: Exception) { diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt index 36e60c6aa9f..aabefda3365 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt @@ -37,10 +37,8 @@ import software.amazon.awssdk.services.codewhispererruntime.model.TargetCode import software.amazon.awssdk.services.codewhispererruntime.model.UserIntent import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger -import software.aws.toolkits.jetbrains.core.awsClient -import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager -import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection import software.aws.toolkits.jetbrains.services.amazonq.codeWhispererUserContext +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContextNew @@ -257,11 +255,8 @@ interface CodeWhispererClientAdaptor { } } -class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispererClientAdaptor { - fun bearerClient(): CodeWhispererRuntimeClient = - ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.getConnectionSettings() - ?.awsClient() - ?: throw Exception("attempt to get bearer client while there is no valid credential") +open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispererClientAdaptor { + fun bearerClient() = QRegionProfileManager.getInstance().getQClient(project) override fun generateCompletionsPaginator(firstRequest: GenerateCompletionsRequest) = sequence { var nextToken: String? = firstRequest.nextToken() @@ -288,7 +283,9 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe // DO NOT directly use this method to fetch customizations, use wrapper [CodeWhispererModelConfigurator.listCustomization()] instead override fun listAvailableCustomizations(): List = - bearerClient().listAvailableCustomizationsPaginator(ListAvailableCustomizationsRequest.builder().build()) + bearerClient().listAvailableCustomizationsPaginator( + ListAvailableCustomizationsRequest.builder().profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn).build() + ) .stream() .toList() .flatMap { resp -> @@ -311,6 +308,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe builder.uploadId(uploadId) builder.targetCodeList(targetCode) builder.userInput(userInput) + builder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) // TODO: client token } @@ -318,6 +316,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe bearerClient().getTestGeneration { builder -> builder.testGenerationJobId(jobId) builder.testGenerationJobGroupName(jobGroupName) + builder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendUserTriggerDecisionTelemetry( @@ -363,6 +362,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } } @@ -409,6 +409,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } } @@ -435,6 +436,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendUserModificationTelemetry( @@ -462,6 +464,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendCodeScanTelemetry( @@ -481,6 +484,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendCodeScanSucceededTelemetry( @@ -503,6 +507,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendCodeScanFailedTelemetry( @@ -522,6 +527,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendCodeFixGenerationTelemetry( @@ -548,6 +554,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendCodeFixAcceptanceTelemetry( @@ -574,6 +581,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendCodeScanRemediationTelemetry( @@ -605,6 +613,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendTestGenerationEvent( @@ -638,10 +647,12 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun listFeatureEvaluations(): ListFeatureEvaluationsResponse = bearerClient().listFeatureEvaluations { it.userContext(codeWhispererUserContext()) + it.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendMetricDataTelemetry(eventName: String, metadata: Map): SendTelemetryEventResponse = @@ -656,6 +667,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendChatAddMessageTelemetry( @@ -694,6 +706,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendChatInteractWithMessageTelemetry( @@ -723,6 +736,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendChatUserModificationTelemetry( @@ -747,6 +761,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } override fun sendInlineChatTelemetry( @@ -782,6 +797,7 @@ class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispe } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) requestBuilder.userContext(codeWhispererUserContext()) + requestBuilder.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } companion object { diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/QStatusBarLoggedInActionGroup.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/QStatusBarLoggedInActionGroup.kt index 4f7833a1fc0..3dd88d45d22 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/QStatusBarLoggedInActionGroup.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/QStatusBarLoggedInActionGroup.kt @@ -8,10 +8,15 @@ import com.intellij.openapi.actionSystem.AnAction import com.intellij.openapi.actionSystem.AnActionEvent import com.intellij.openapi.actionSystem.DefaultActionGroup import com.intellij.openapi.actionSystem.Separator +import com.intellij.openapi.project.Project import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.actions.SsoLogoutAction import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection +import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection +import software.aws.toolkits.jetbrains.core.credentials.sono.isSono +import software.aws.toolkits.jetbrains.services.amazonq.actions.QSwitchProfilesAction +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererConnectOnGithubAction import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererLearnMoreAction import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererProvideFeedbackAction @@ -48,18 +53,11 @@ class QStatusBarLoggedInActionGroup : DefaultActionGroup() { } override fun getChildren(e: AnActionEvent?) = e?.project?.let { + val isPendingActiveProfile = QRegionProfileManager.getInstance().hasValidConnectionButNoActiveProfile(it) buildList { - add(Separator.create()) - add(Separator.create(message("codewhisperer.statusbar.sub_menu.inline.title"))) - addAll(buildActionListForInlineSuggestions(it, actionProvider)) - - add(Separator.create()) - add(Separator.create(message("codewhisperer.statusbar.sub_menu.security_scans.title"))) - addAll(buildActionListForCodeScan(it, actionProvider)) - - add(Separator.create()) - add(Separator.create(message("codewhisperer.statusbar.sub_menu.other_features.title"))) - addAll(buildActionListForOtherFeatures(it, actionProvider)) + if (!isPendingActiveProfile) { + addAll(buildActionListForActiveProfileSelected(it, actionProvider)) + } add(Separator.create()) add(Separator.create(message("codewhisperer.statusbar.sub_menu.connect_help.title"))) @@ -67,6 +65,10 @@ class QStatusBarLoggedInActionGroup : DefaultActionGroup() { add(Separator.create()) add(CodeWhispererShowSettingsAction()) + ( + ToolkitConnectionManager.getInstance(it).activeConnectionForFeature(QConnection.getInstance()) as? AwsBearerTokenConnection + )?.takeIf { !it.isSono() } + ?.let { add(QSwitchProfilesAction()) } ToolkitConnectionManager.getInstance(it).activeConnectionForFeature(CodeWhispererConnection.getInstance())?.let { c -> (c as? AwsBearerTokenConnection)?.let { connection -> add(SsoLogoutAction(connection)) @@ -74,4 +76,21 @@ class QStatusBarLoggedInActionGroup : DefaultActionGroup() { } }.toTypedArray() }.orEmpty() + + private fun buildActionListForActiveProfileSelected( + project: Project, + actionProvider: ActionProvider, + ): List = buildList { + add(Separator.create()) + add(Separator.create(message("codewhisperer.statusbar.sub_menu.inline.title"))) + addAll(buildActionListForInlineSuggestions(project, actionProvider)) + + add(Separator.create()) + add(Separator.create(message("codewhisperer.statusbar.sub_menu.security_scans.title"))) + addAll(buildActionListForCodeScan(project, actionProvider)) + + add(Separator.create()) + add(Separator.create(message("codewhisperer.statusbar.sub_menu.other_features.title"))) + addAll(buildActionListForOtherFeatures(project, actionProvider)) + } } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererAutoTriggerHandler.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererAutoTriggerHandler.kt index fe73ff91e48..73b9f59be6f 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererAutoTriggerHandler.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererAutoTriggerHandler.kt @@ -7,6 +7,7 @@ import com.intellij.openapi.editor.Editor import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.codewhisperer.model.LatencyContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.TriggerTypeInfo import software.aws.toolkits.telemetry.CodewhispererTriggerType @@ -17,6 +18,10 @@ interface CodeWhispererAutoTriggerHandler { automatedTriggerType: CodeWhispererAutomatedTriggerType, latencyContext: LatencyContext, ) { + val project = editor.project ?: return + if (QRegionProfileManager.getInstance().hasValidConnectionButNoActiveProfile(project)) { + return + } val triggerTypeInfo = TriggerTypeInfo(CodewhispererTriggerType.AutoTrigger, automatedTriggerType) LOG.debug { "autotriggering CodeWhisperer with type ${automatedTriggerType.telemetryType}" } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt index 05d814b94a9..50a579fedd1 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt @@ -56,6 +56,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.SUPPLEMENTAL_CONTEXT_TIM import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.GetConfigurationFromServerParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerConfigurations +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorManager @@ -238,7 +239,8 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { requestContext.fileContextInfo, requestContext.awaitSupplementalContext(), requestContext.customizationArn, - requestContext.workspaceId + requestContext.profileArn, + requestContext.workspaceId, ) ) @@ -673,6 +675,8 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { // 5. customization val customizationArn = CodeWhispererModelConfigurator.getInstance().activeCustomization(project)?.arn + val profileArn = QRegionProfileManager.getInstance().activeProfile(project)?.arn + var workspaceId: String? = null try { val workspacesInfos = getWorkspaceIds(project).get().workspaces @@ -688,8 +692,17 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { LOG.warn { "Cannot get workspaceId from LSP'$e'" } } return RequestContext( - project, editor, triggerTypeInfo, caretPosition, - fileContext, supplementalContext, connection, latencyContext, customizationArn, workspaceId + project, + editor, + triggerTypeInfo, + caretPosition, + fileContext, + supplementalContext, + connection, + latencyContext, + customizationArn, + profileArn, + workspaceId, ) } @@ -833,6 +846,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { fileContextInfo: FileContextInfo, supplementalContext: SupplementalContextInfo?, customizationArn: String?, + profileArn: String?, workspaceId: String?, ): GenerateCompletionsRequest { val programmingLanguage = ProgrammingLanguage.builder() @@ -862,6 +876,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { .referenceTrackerConfiguration { it.recommendationsWithReferences(includeCodeWithReference) } .customizationArn(customizationArn) .optOutPreference(getTelemetryOptOutPreference()) + .profileArn(profileArn) .workspaceId(workspaceId) .build() } @@ -878,6 +893,7 @@ data class RequestContext( val connection: ToolkitConnection?, val latencyContext: LatencyContext, val customizationArn: String?, + val profileArn: String?, val workspaceId: String?, ) { // TODO: should make the entire getRequestContext() suspend function instead of making supplemental context only diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererServiceNew.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererServiceNew.kt index c79ff7c0123..633a20ac0f9 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererServiceNew.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererServiceNew.kt @@ -53,6 +53,7 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection import software.aws.toolkits.jetbrains.services.amazonq.SUPPLEMENTAL_CONTEXT_TIMEOUT +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorManagerNew @@ -243,7 +244,8 @@ class CodeWhispererServiceNew(private val cs: CoroutineScope) : Disposable { buildCodeWhispererRequest( requestContext.fileContextInfo, requestContext.awaitSupplementalContext(), - requestContext.customizationArn + requestContext.customizationArn, + requestContext.profileArn ) ) @@ -707,7 +709,9 @@ class CodeWhispererServiceNew(private val cs: CoroutineScope) : Disposable { // 5. customization val customizationArn = CodeWhispererModelConfigurator.getInstance().activeCustomization(project)?.arn - return RequestContextNew(project, editor, triggerTypeInfo, caretPosition, fileContext, supplementalContext, connection, customizationArn) + val profileArn = QRegionProfileManager.getInstance().activeProfile(project)?.arn + + return RequestContextNew(project, editor, triggerTypeInfo, caretPosition, fileContext, supplementalContext, connection, customizationArn, profileArn) } fun validateResponse(response: GenerateCompletionsResponse): GenerateCompletionsResponse { @@ -825,6 +829,7 @@ class CodeWhispererServiceNew(private val cs: CoroutineScope) : Disposable { fileContextInfo: FileContextInfo, supplementalContext: SupplementalContextInfo?, customizationArn: String?, + profileArn: String?, ): GenerateCompletionsRequest { val programmingLanguage = ProgrammingLanguage.builder() .languageName(fileContextInfo.programmingLanguage.toCodeWhispererRuntimeLanguage().languageId) @@ -853,6 +858,7 @@ class CodeWhispererServiceNew(private val cs: CoroutineScope) : Disposable { .referenceTrackerConfiguration { it.recommendationsWithReferences(includeCodeWithReference) } .customizationArn(customizationArn) .optOutPreference(getTelemetryOptOutPreference()) + .profileArn(profileArn) .build() } } @@ -867,6 +873,7 @@ data class RequestContextNew( private val supplementalContextDeferred: Deferred, val connection: ToolkitConnection?, val customizationArn: String?, + val profileArn: String?, ) { // TODO: should make the entire getRequestContext() suspend function instead of making supplemental context only var supplementalContext: SupplementalContextInfo? = null diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt index 1e36ceddc5d..9e328c54b81 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt @@ -6,7 +6,6 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.util import com.intellij.openapi.actionSystem.DataKey import com.intellij.openapi.editor.markup.EffectType import com.intellij.openapi.editor.markup.TextAttributes -import com.intellij.openapi.util.registry.Registry import com.intellij.ui.JBColor import software.amazon.awssdk.regions.Region import software.amazon.awssdk.services.codewhispererruntime.model.AccessDeniedException @@ -154,12 +153,8 @@ object CodeWhispererConstants { } object Config { - val CODEWHISPERER_ENDPOINT - get() = System.getenv("__CODEWHISPERER_ENDPOINT") ?: Registry.get("amazon.q.endpoint").asString() - const val CODEWHISPERER_IDPOOL_ID = "us-east-1:70717e99-906f-4add-908c-bd9074a2f5b9" val Sigv4ClientRegion = Region.US_EAST_1 - val BearerClientRegion = Region.US_EAST_1 } object Customization { diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererEndpointCustomizer.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererEndpointCustomizer.kt index 004181a6ec5..8d74ecca7e0 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererEndpointCustomizer.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererEndpointCustomizer.kt @@ -17,9 +17,12 @@ import software.amazon.awssdk.core.retry.RetryPolicy import software.amazon.awssdk.http.SdkHttpRequest import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient import software.amazon.awssdk.http.nio.netty.ProxyConfiguration +import software.amazon.awssdk.regions.Region import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClientBuilder import software.amazon.awssdk.services.codewhispererstreaming.CodeWhispererStreamingAsyncClientBuilder import software.aws.toolkits.core.ToolkitClientCustomizer +import software.aws.toolkits.core.utils.tryOrNull +import software.aws.toolkits.jetbrains.services.amazonq.profile.QEndpoints import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings import java.net.Proxy import java.net.URI @@ -37,10 +40,12 @@ class CodeWhispererEndpointCustomizer : ToolkitClientCustomizer { clientOverrideConfiguration: ClientOverrideConfiguration.Builder, ) { if (builder is CodeWhispererRuntimeClientBuilder || builder is CodeWhispererStreamingAsyncClientBuilder) { - val endpoint = URI.create(CodeWhispererConstants.Config.CODEWHISPERER_ENDPOINT) + val endpoint = tryOrNull { QEndpoints.getQEndpointWithRegion(regionId) } + ?.let { URI.create(it) } + ?: URI.create(QEndpoints.Q_DEFAULT_SERVICE_CONFIG.ENDPOINT) builder .endpointOverride(endpoint) - .region(CodeWhispererConstants.Config.BearerClientRegion) + .region(Region.of(regionId)) clientOverrideConfiguration.retryPolicy(RetryPolicy.none()) clientOverrideConfiguration.addExecutionInterceptor( object : ExecutionInterceptor { diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererZipUploadManager.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererZipUploadManager.kt index 8d997645a14..f964092a8d5 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererZipUploadManager.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererZipUploadManager.kt @@ -22,6 +22,7 @@ import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.jetbrains.core.AwsClientManager import software.aws.toolkits.jetbrains.services.amazonq.RetryableOperation +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanServerException import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanSession.Companion.APPLICATION_ZIP import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanSession.Companion.AWS_KMS @@ -202,6 +203,7 @@ class CodeWhispererZipUploadManager(private val project: Project) { UploadContext.fromCodeAnalysisUploadContext(CodeAnalysisUploadContext.builder().codeScanName(taskName).build()) } ) + .profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) .build() ) }, diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt index 1e81a695308..bf3db6cdf29 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt @@ -177,7 +177,8 @@ internal class CodeWhispererCodeCoverageTrackerTestPython : CodeWhispererCodeCov null, mock(), aString(), - null + aString(), + aString(), ) val responseContext = ResponseContext("sessionId") val recommendationContext = RecommendationContext( diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt index 56f6f66b517..378e2703756 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt @@ -3,6 +3,7 @@ package software.aws.toolkits.jetbrains.services.codewhisperer +import com.intellij.openapi.project.Project import com.intellij.testFramework.ApplicationRule import com.intellij.testFramework.DisposableRule import com.intellij.testFramework.ProjectRule @@ -34,6 +35,7 @@ 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.SONO_URL import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import kotlin.reflect.full.memberFunctions import kotlin.test.Test @@ -62,7 +64,7 @@ class CodeWhispererFeatureConfigServiceTest { @Test fun `test highlightCommand returns non-empty`() { - mockClientManagerRule.create().stub { + val mockClient = mockClientManagerRule.create().stub { on { listFeatureEvaluations(any()) } doReturn ListFeatureEvaluationsResponse.builder().featureEvaluations( listOf( FeatureEvaluation.builder() @@ -74,9 +76,16 @@ class CodeWhispererFeatureConfigServiceTest { ).build() } + projectRule.project.replaceService( + QRegionProfileManager::class.java, + mock { on { getQClient(any(), eq(CodeWhispererRuntimeClient::class)) } doReturn mockClient }, + disposableRule.disposable + ) + val mockTokenSettings = mock { on { providerId } doReturn "mock" on { region } doReturn AwsRegion.GLOBAL + on { withRegion(any()) } doReturn this.mock } val mockSsoConnection = mock { diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererServiceTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererServiceTest.kt index add4a9909cd..4ac3f45eddd 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererServiceTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererServiceTest.kt @@ -213,7 +213,8 @@ class CodeWhispererServiceTest { connection = ToolkitConnectionManager.getInstance(projectRule.project).activeConnection(), latencyContext = LatencyContext(), customizationArn = "fake-arn", - workspaceId = null + profileArn = "fake-arn", + workspaceId = null, ) ) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt index a3e97f60c38..e000bbaeafc 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt @@ -37,6 +37,7 @@ import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRul import software.aws.toolkits.jetbrains.core.credentials.MockToolkitAuthManagerRule import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.codeWhispererRecommendationActionId import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonFileName import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponse @@ -88,6 +89,7 @@ open class CodeWhispererTestBase { protected lateinit var settingsManager: CodeWhispererSettings private lateinit var originalExplorerActionState: CodeWhispererExploreActionState private lateinit var originalSettings: CodeWhispererConfiguration + private lateinit var qRegionProfileManagerSpy: QRegionProfileManager @Before open fun setUp() { @@ -169,6 +171,16 @@ open class CodeWhispererTestBase { val conn = authManagerRule.createConnection(ManagedSsoProfile("us-east-1", "url", Q_SCOPES)) ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(conn) + + qRegionProfileManagerSpy = spy(QRegionProfileManager.getInstance()) + qRegionProfileManagerSpy.stub { + onGeneric { + hasValidConnectionButNoActiveProfile(any()) + } doAnswer { + false + } + } + ApplicationManager.getApplication().replaceService(QRegionProfileManager::class.java, qRegionProfileManagerSpy, disposableRule.disposable) } @After diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QEndpointsTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QEndpointsTest.kt new file mode 100644 index 00000000000..057c61030f7 --- /dev/null +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QEndpointsTest.kt @@ -0,0 +1,41 @@ +// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package software.aws.toolkits.jetbrains.services.codewhisperer + +import com.intellij.testFramework.ApplicationExtension +import com.intellij.testFramework.fixtures.BasePlatformTestCase +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.extension.ExtendWith +import org.junit.jupiter.api.extension.RegisterExtension +import software.aws.toolkits.jetbrains.services.amazonq.profile.QEndpoints +import software.aws.toolkits.jetbrains.utils.rules.RegistryExtension + +@ExtendWith(ApplicationExtension::class) +class QEndpointsTest : BasePlatformTestCase() { + + @JvmField + @RegisterExtension + val registryExtension = RegistryExtension() + + @Test + fun `test default registry value and parse`() { + val testJson = """ + [ + {"region": "us-east-1", "endpoint": "https://codewhisperer.us-east-1.amazonaws.com/"}, + {"region": "eu-central-1", "endpoint": "https://rts.prod-eu-central-1.codewhisperer.ai.aws.dev/"} + ] + """.trimIndent() + + registryExtension.setValue("amazon.q.endpoints.json", testJson) + + val parsed = QEndpoints.listRegionEndpoints() + assertEquals(2, parsed.size) + + val iad = parsed.first { it.region == "us-east-1" } + assertEquals("https://codewhisperer.us-east-1.amazonaws.com/", iad.endpoint) + + val fra = parsed.first { it.region == "eu-central-1" } + assertEquals("https://rts.prod-eu-central-1.codewhisperer.ai.aws.dev/", fra.endpoint) + } +} diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt new file mode 100644 index 00000000000..640d77e7cd1 --- /dev/null +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt @@ -0,0 +1,387 @@ +// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +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.util.xmlb.XmlSerializer +import org.assertj.core.api.Assertions.assertThat +import org.jdom.output.XMLOutputter +import org.junit.Before +import org.junit.Rule +import org.junit.Test +import org.mockito.kotlin.any +import org.mockito.kotlin.doReturn +import org.mockito.kotlin.mock +import org.mockito.kotlin.stub +import software.amazon.awssdk.core.pagination.sync.SdkIterable +import software.amazon.awssdk.regions.Region +import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient +import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableProfilesRequest +import software.amazon.awssdk.services.codewhispererruntime.model.Profile +import software.amazon.awssdk.services.codewhispererruntime.paginators.ListAvailableProfilesIterable +import software.amazon.awssdk.services.ssooidc.SsoOidcClient +import software.aws.toolkits.core.region.AwsRegion +import software.aws.toolkits.jetbrains.core.MockClientManager +import software.aws.toolkits.jetbrains.core.MockClientManagerRule +import software.aws.toolkits.jetbrains.core.MockResourceCacheRule +import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection +import software.aws.toolkits.jetbrains.core.credentials.ManagedSsoProfile +import software.aws.toolkits.jetbrains.core.credentials.MockToolkitAuthManagerRule +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.region.MockRegionProviderRule +import software.aws.toolkits.jetbrains.services.amazonq.profile.QEndpoints +import software.aws.toolkits.jetbrains.services.amazonq.profile.QProfileResources +import software.aws.toolkits.jetbrains.services.amazonq.profile.QProfileState +import software.aws.toolkits.jetbrains.services.amazonq.profile.QProfileSwitchIntent +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener +import software.aws.toolkits.jetbrains.utils.satisfiesKt +import software.aws.toolkits.jetbrains.utils.xmlElement +import java.net.URI +import java.util.function.Consumer +import kotlin.test.fail + +// TODO: should use junit5 +class QRegionProfileManagerTest { + @Rule + @JvmField + val projectRule = ProjectRule() + + @Rule + @JvmField + val authRule = MockToolkitAuthManagerRule() + + @JvmField + @Rule + val clientRule = MockClientManagerRule() + + @Rule + @JvmField + val regionProviderRule = MockRegionProviderRule() + + @JvmField + @Rule + val disposableRule = DisposableRule() + + @get:Rule + val resourceCache = MockResourceCacheRule() + + private lateinit var sut: QRegionProfileManager + private val project: Project + get() = projectRule.project + + @Before + fun setup() { + clientRule.create() + regionProviderRule.addRegion(AwsRegion("us-east-1", "US East (N. Virginia)", "aws")) + regionProviderRule.addRegion(AwsRegion("eu-central-1", "Europe (Frankfurt)", "aws")) + sut = QRegionProfileManager() + val conn = authRule.createConnection(ManagedSsoProfile(ssoRegion = "us-east-1", startUrl = "", scopes = Q_SCOPES)) + ToolkitConnectionManager.getInstance(project).switchConnection(conn) + } + + @Test + fun `switchProfile should switch the current connection(project) to the selected profile`() { + sut.switchProfile(project, QRegionProfile(arn = "arn", profileName = "foo_profile"), QProfileSwitchIntent.User) + assertThat(sut.activeProfile(project)).isEqualTo(QRegionProfile(arn = "arn", profileName = "foo_profile")) + + sut.switchProfile(project, QRegionProfile(arn = "another_arn", profileName = "bar_profile"), QProfileSwitchIntent.User) + assertThat(sut.activeProfile(project)).isEqualTo(QRegionProfile(arn = "another_arn", profileName = "bar_profile")) + } + + @Test + fun `switchProfile should return null if user is not connected`() { + sut.switchProfile(project, QRegionProfile(arn = "arn", profileName = "foo_profile"), QProfileSwitchIntent.User) + assertThat(sut.activeProfile(project)).isEqualTo(QRegionProfile(arn = "arn", profileName = "foo_profile")) + + ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let { + if (it is AwsBearerTokenConnection) { + logoutFromSsoConnection(project, it) + } + } + + assertThat(ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())).isNull() + assertThat(sut.activeProfile(project)).isNull() + } + + @Test + fun `data is cleared when user logs out`() { + sut.switchProfile(project, QRegionProfile(arn = "arn", profileName = "foo_profile"), QProfileSwitchIntent.User) + assertThat(sut.activeProfile(project)).isEqualTo(QRegionProfile(arn = "arn", profileName = "foo_profile")) + + ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let { + if (it is AwsBearerTokenConnection) { + logoutFromSsoConnection(project, it) + } + } + + assertThat(sut.state).satisfiesKt { + assertThat(it.connectionIdToActiveProfile).isEmpty() + assertThat(it.connectionIdToProfileList).isEmpty() + } + } + + @Test + fun `switch should send message onProfileChanged for active switch`() { + var cnt = 0 + project.messageBus.connect(disposableRule.disposable).subscribe( + QRegionProfileSelectedListener.TOPIC, + object : QRegionProfileSelectedListener { + override fun onProfileSelected(project: Project, profile: QRegionProfile?) { + cnt += 1 + } + } + ) + + assertThat(cnt).isEqualTo(0) + sut.switchProfile(project, QRegionProfile(arn = "arn", profileName = "foo_profile"), QProfileSwitchIntent.Reload) + assertThat(cnt).isEqualTo(1) + sut.switchProfile(project, QRegionProfile(arn = "another_arn", profileName = "BAR_PROFILE"), QProfileSwitchIntent.Reload) + assertThat(cnt).isEqualTo(2) + } + + @Test + fun `listProfiles will call each client to get profiles`() { + val client = clientRule.create() + val mockResponse: SdkIterable = SdkIterable { + listOf( + Profile.builder().profileName("FOO").arn("foo").build(), + ).toMutableList().iterator() + } + + val mockResponse2: SdkIterable = SdkIterable { + listOf( + Profile.builder().profileName("BAR").arn("bar").build(), + ).toMutableList().iterator() + } + + val iterable: ListAvailableProfilesIterable = mock { + on { it.profiles() } doReturn mockResponse doReturn mockResponse2 + } + + // TODO: not sure if we can mock client with different region different response? + client.stub { + onGeneric { listAvailableProfilesPaginator(any>()) } doReturn iterable + } + val connectionSettings = sut.getQClientSettings(project) + resourceCache.addEntry(connectionSettings, QProfileResources.LIST_REGION_PROFILES, QProfileResources.LIST_REGION_PROFILES.fetch(connectionSettings)) + + assertThat(sut.listRegionProfiles(project)) + .hasSize(2) + .containsExactlyInAnyOrder( + QRegionProfile("FOO", "foo"), + QRegionProfile("BAR", "bar") + ) + } + + @Test + fun `validateProfile should cross validate selected profile with latest API response for current project and remove it if its not longer accessible`() { + val client = clientRule.create() + val mockResponse: SdkIterable = SdkIterable { + listOf( + Profile.builder().profileName("foo").arn("foo-arn-v2").build(), + Profile.builder().profileName("bar").arn("bar-arn").build(), + ).toMutableList().iterator() + } + val iterable: ListAvailableProfilesIterable = mock { + on { it.profiles() } doReturn mockResponse + } + client.stub { + onGeneric { listAvailableProfilesPaginator(any>()) } doReturn iterable + } + + val activeConn = + ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) ?: fail("connection shouldn't be null") + val anotherConn = authRule.createConnection(ManagedSsoProfile(ssoRegion = "us-east-1", startUrl = "anotherUrl", scopes = Q_SCOPES)) + val fooProfile = QRegionProfile("foo", "foo-arn") + val barProfile = QRegionProfile("bar", "bar-arn") + val state = QProfileState().apply { + this.connectionIdToActiveProfile[activeConn.id] = fooProfile + this.connectionIdToActiveProfile[anotherConn.id] = barProfile + } + sut.loadState(state) + assertThat(sut.activeProfile(project)).isEqualTo(fooProfile) + + sut.validateProfile(project) + assertThat(sut.activeProfile(project)).isNull() + assertThat(sut.state.connectionIdToActiveProfile).isEqualTo(mapOf(anotherConn.id to barProfile)) + } + + @Test + fun `clientSettings should return the region Q profile specify`() { + MockClientManager.useRealImplementations(disposableRule.disposable) + sut.switchProfile( + project, + QRegionProfile(arn = "arn:aws:codewhisperer:eu-central-1:123456789012:profile/FOO_PROFILE", profileName = "FOO_PROFILE"), + QProfileSwitchIntent.User + ) + assertThat( + sut.activeProfile(project) + ).isEqualTo(QRegionProfile(arn = "arn:aws:codewhisperer:eu-central-1:123456789012:profile/FOO_PROFILE", profileName = "FOO_PROFILE")) + + val settings = sut.getQClientSettings(project) + assertThat(settings.region.id).isEqualTo(Region.EU_CENTRAL_1.id()) + + sut.switchProfile( + project, + QRegionProfile(arn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/BAR_PROFILE", profileName = "BAR_PROFILE"), + QProfileSwitchIntent.User + ) + assertThat( + sut.activeProfile(project) + ).isEqualTo(QRegionProfile(arn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/BAR_PROFILE", profileName = "BAR_PROFILE")) + + val settings2 = sut.getQClientSettings(project) + assertThat(settings2.region.id).isEqualTo(Region.US_EAST_1.id()) + } + + @Test + fun `getClient should return correct client with region and endpoint`() { + MockClientManager.useRealImplementations(disposableRule.disposable) + + sut.switchProfile( + project, + QRegionProfile(arn = "arn:aws:codewhisperer:eu-central-1:123456789012:profile/FOO_PROFILE", profileName = "FOO_PROFILE"), + QProfileSwitchIntent.User + ) + assertThat( + sut.activeProfile(project) + ).isEqualTo(QRegionProfile(arn = "arn:aws:codewhisperer:eu-central-1:123456789012:profile/FOO_PROFILE", profileName = "FOO_PROFILE")) + assertThat(sut.getQClientSettings(project).region.id).isEqualTo(Region.EU_CENTRAL_1.id()) + + val client = sut.getQClient(project) + assertThat(client).isInstanceOf(CodeWhispererRuntimeClient::class.java) + assertThat(client.serviceClientConfiguration().region()).isEqualTo(Region.EU_CENTRAL_1) + assertThat( + client.serviceClientConfiguration().endpointOverride().get() + ).isEqualTo(URI.create(QEndpoints.getQEndpointWithRegion(Region.EU_CENTRAL_1.id()))) + + sut.switchProfile( + project, + QRegionProfile(arn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/BAR_PROFILE", profileName = "BAR_PROFILE"), + QProfileSwitchIntent.User + ) + assertThat( + sut.activeProfile(project) + ).isEqualTo(QRegionProfile(arn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/BAR_PROFILE", profileName = "BAR_PROFILE")) + assertThat(sut.getQClientSettings(project).region.id).isEqualTo(Region.US_EAST_1.id()) + + val client2 = sut.getQClient(project) + assertThat(client2).isInstanceOf(CodeWhispererRuntimeClient::class.java) + assertThat(client2.serviceClientConfiguration().region()).isEqualTo(Region.US_EAST_1) + assertThat( + client2.serviceClientConfiguration().endpointOverride().get() + ).isEqualTo(URI.create(QEndpoints.getQEndpointWithRegion(Region.US_EAST_1.id()))) + } + + @Test + fun `deserialize empty data`() { + val element = xmlElement( + """ + + + """ + ) + val actual = XmlSerializer.deserialize(element, QProfileState::class.java) + assertThat(actual.connectionIdToActiveProfile).hasSize(0) + assertThat(actual.connectionIdToProfileList).hasSize(0) + } + + @Test + fun `serialize with data`() { + val element = xmlElement( + """ + + + """.trimIndent() + ) + + val state = QProfileState().apply { + this.connectionIdToActiveProfile.putAll( + mapOf( + "conn-123" to QRegionProfile( + profileName = "myActiveProfile", arn = "arn:aws:codewhisperer:us-west-2:123456789012:profile/myActiveProfile" + ) + ) + ) + + connectionIdToProfileList.putAll( + mapOf("conn-123" to 2) + ) + } + + XmlSerializer.serializeInto(state, element) + val actualXmlString = XMLOutputter().outputString(element) + val expectedXmlString = + "\n" + + "" + + "" + + "" + + assertThat(actualXmlString).isEqualTo(expectedXmlString) + } + + @Test + fun `deserialize with data`() { + val element = xmlElement( + """ + + + + + """.trimIndent() + ) + + val actualState = XmlSerializer.deserialize(element, QProfileState::class.java) + + assertThat(actualState.connectionIdToActiveProfile).hasSize(1) + val activeProfile = actualState.connectionIdToActiveProfile["conn-123"] + assertThat(activeProfile).isEqualTo( + QRegionProfile( + profileName = "myActiveProfile", + arn = "arn:aws:codewhisperer:us-west-2:123456789012:profile/myActiveProfile" + ) + ) + + assertThat(actualState.connectionIdToProfileList).hasSize(1) + val profileList = actualState.connectionIdToProfileList["conn-123"] + assertThat(profileList).isEqualTo(2) + } +} diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestUtil.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestUtil.kt index 3bf772e86d5..1178a7ff967 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestUtil.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestUtil.kt @@ -258,7 +258,8 @@ fun aRequestContext( aString() ), customizationArn = null, - workspaceId = null + profileArn = null, + workspaceId = null, ) } diff --git a/plugins/amazonq/mynah-ui/src/mynah-ui/ui/main.ts b/plugins/amazonq/mynah-ui/src/mynah-ui/ui/main.ts index 1be8f151a72..1d033f1cc65 100644 --- a/plugins/amazonq/mynah-ui/src/mynah-ui/ui/main.ts +++ b/plugins/amazonq/mynah-ui/src/mynah-ui/ui/main.ts @@ -42,6 +42,7 @@ export const createMynahUI = ( codeScanEnabled: boolean, codeTestEnabled: boolean, highlightCommand?: QuickActionCommand, + profileName?: string ) => { let disclaimerCardActive = !disclaimerAcknowledged @@ -88,7 +89,8 @@ export const createMynahUI = ( isDocEnabled, isCodeScanEnabled, isCodeTestEnabled, - highlightCommand + highlightCommand, + profileName }) // eslint-disable-next-line prefer-const diff --git a/plugins/amazonq/mynah-ui/src/mynah-ui/ui/tabs/generator.ts b/plugins/amazonq/mynah-ui/src/mynah-ui/ui/tabs/generator.ts index 63682d596e7..f94eb193c98 100644 --- a/plugins/amazonq/mynah-ui/src/mynah-ui/ui/tabs/generator.ts +++ b/plugins/amazonq/mynah-ui/src/mynah-ui/ui/tabs/generator.ts @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { ChatItemType, MynahUIDataModel, QuickActionCommandGroup, QuickActionCommand } from '@aws/mynah-ui-chat' +import {ChatItemType, MynahUIDataModel, QuickActionCommandGroup, QuickActionCommand, ChatItem} from '@aws/mynah-ui-chat' import { TabType } from '../storages/tabsStorage' import { FollowUpGenerator } from '../followUps/generator' import { QuickActionGenerator } from '../quickActions/generator' @@ -16,12 +16,14 @@ export interface TabDataGeneratorProps { isCodeScanEnabled: boolean isCodeTestEnabled: boolean highlightCommand?: QuickActionCommand + profileName?: string } export class TabDataGenerator { private followUpsGenerator: FollowUpGenerator public quickActionsGenerator: QuickActionGenerator public highlightCommand?: QuickActionCommand + profileName?: string private tabTitle: Map = new Map([ ['unknown', 'Chat'], @@ -91,6 +93,20 @@ What would you like to work on?`, isCodeTestEnabled: props.isCodeTestEnabled, }) this.highlightCommand = props.highlightCommand + this.profileName = props.profileName + } + + private get regionProfileCard(): ChatItem | undefined { + console.log('[DEBUG] Received profileName:', this.profileName) + if (!this.profileName) { + return undefined + } + return { + type: ChatItemType.ANSWER, + body: `You are using the ${this.profileName} profile for this chat period`, + status: 'info', + messageId: 'regionProfile', + } } public getTabData(tabType: TabType, needWelcomeMessages: boolean, taskName?: string): MynahUIDataModel { @@ -103,7 +119,8 @@ What would you like to work on?`, contextCommands: this.getContextCommands(tabType), chatItems: needWelcomeMessages ? [ - { + ...(this.regionProfileCard ? [this.regionProfileCard] : []), + { type: ChatItemType.ANSWER, body: this.tabWelcomeMessage.get(tabType), }, @@ -112,7 +129,7 @@ What would you like to work on?`, followUp: this.followUpsGenerator.generateWelcomeBlockForTab(tabType), }, ] - : [], + : [...(this.regionProfileCard ? [this.regionProfileCard] : [])], } } diff --git a/plugins/amazonq/mynah-ui/src/mynah-ui/ui/walkthrough/welcome.ts b/plugins/amazonq/mynah-ui/src/mynah-ui/ui/walkthrough/welcome.ts index e5fe67fdd7b..feb879d34f5 100644 --- a/plugins/amazonq/mynah-ui/src/mynah-ui/ui/walkthrough/welcome.ts +++ b/plugins/amazonq/mynah-ui/src/mynah-ui/ui/walkthrough/welcome.ts @@ -14,6 +14,14 @@ export const welcomeScreenTabData = (tabs: TabDataGenerator): MynahUITabStoreTab tabTitle: 'Welcome to Q', tabBackground: true, chatItems: [ + ...(tabs.profileName + ? [{ + type: ChatItemType.ANSWER, + icon: MynahIcons.INFO, + messageId: 'profile-info', + body: `You're using the ${tabs.profileName} profile for this chat period.`, + }] + : []), { type: ChatItemType.ANSWER, icon: MynahIcons.ASTERISK, diff --git a/plugins/amazonq/shared/jetbrains-community/resources/software/aws/toolkits/resources/AmazonQBundle.properties b/plugins/amazonq/shared/jetbrains-community/resources/software/aws/toolkits/resources/AmazonQBundle.properties index fab9733aa8d..79ecca81754 100644 --- a/plugins/amazonq/shared/jetbrains-community/resources/software/aws/toolkits/resources/AmazonQBundle.properties +++ b/plugins/amazonq/shared/jetbrains-community/resources/software/aws/toolkits/resources/AmazonQBundle.properties @@ -10,4 +10,13 @@ amazonqInlineChat.popup.title=Enter Instructions for Q amazonq.refresh.panel=Refresh Chat Session amazonq.title=Amazon Q amazonq.workspace.settings.open.prompt=Workspace index is now enabled. You can disable it from Amazon Q settings. -q.hello=Hello +action.q.profile.usage.text=You changed profile +action.q.profile.usage=You're using the '{0}' profile for Amazon Q. +action.q.switchProfiles.text=Change profile +action.q.switchProfiles.dialog.text=Amazon Q Developer Profile +action.q.switchProfiles.dialog.account.label=Account: {0} +action.q.switchProfiles.dialog.panel.text=Change your Q Developer profile +action.q.switchProfiles.dialog.panel.description=Choose the profile that meets your current working needs. +action.q.switchProfiles.dialog.panel.warning=When you change profiles, you will no longer have access to your current customizations, chats, code reviews, or any other code or content being generated by Amazon Q. +general.ok=OK +general.cancel=Cancel diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt index 667e8e8f844..d77663430fb 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt @@ -12,9 +12,9 @@ import software.amazon.awssdk.services.codewhispererruntime.model.FeatureValue import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.error import software.aws.toolkits.core.utils.getLogger -import software.aws.toolkits.jetbrains.core.awsClient import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.utils.isQExpired @Service @@ -32,8 +32,9 @@ class CodeWhispererFeatureConfigService { LOG.debug { "Fetching feature configs" } try { - val response = connection.getConnectionSettings().awsClient().listFeatureEvaluations { + val response = QRegionProfileManager.getInstance().getQClient(project).listFeatureEvaluations { it.userContext(codeWhispererUserContext()) + it.profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn) } ?: return // Simply force overwrite feature configs from server response, no needed to check existing values. @@ -113,7 +114,8 @@ class CodeWhispererFeatureConfigService { val availableCustomizations = calculateIfIamIdentityCenterConnection(project) { try { - connection.getConnectionSettings().awsClient().listAvailableCustomizationsPaginator {} + QRegionProfileManager.getInstance().getQClient(project) + .listAvailableCustomizationsPaginator {} .flatMap { resp -> resp.customizations().map { it.arn() diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/actions/QSwitchProfilesAction.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/actions/QSwitchProfilesAction.kt new file mode 100644 index 00000000000..2e8949d0f66 --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/actions/QSwitchProfilesAction.kt @@ -0,0 +1,58 @@ +// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package software.aws.toolkits.jetbrains.services.amazonq.actions + +import com.intellij.icons.AllIcons +import com.intellij.openapi.actionSystem.ActionUpdateThread +import com.intellij.openapi.actionSystem.AnAction +import com.intellij.openapi.actionSystem.AnActionEvent +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.project.DumbAware +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.services.amazonq.profile.QProfileSwitchIntent +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileDialog +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager +import software.aws.toolkits.resources.AmazonQBundle.message +import software.aws.toolkits.telemetry.MetricResult +import software.aws.toolkits.telemetry.Telemetry + +class QSwitchProfilesAction : AnAction(message("action.q.switchProfiles.text")), DumbAware { + + override fun getActionUpdateThread() = ActionUpdateThread.BGT + + override fun update(e: AnActionEvent) { + e.presentation.icon = AllIcons.Actions.SwapPanels + } + + override fun actionPerformed(e: AnActionEvent) { + val project = e.project ?: return + ApplicationManager.getApplication().executeOnPooledThread { + val profiles = try { + QRegionProfileManager.getInstance().listRegionProfiles(project) + } catch (e: Exception) { + val conn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) as? AwsBearerTokenConnection + Telemetry.amazonq.didSelectProfile.use { span -> + span.source(QProfileSwitchIntent.User.value) + .amazonQProfileRegion(QRegionProfileManager.getInstance().activeProfile(project)?.region ?: "not-set") + .ssoRegion(conn?.region) + .credentialStartUrl(conn?.startUrl) + .result(MetricResult.Failed) + .reason(e.message) + } + throw e + } + ?: error("Attempted to fetch profiles while there does not exist") + val selectedProfile = QRegionProfileManager.getInstance().activeProfile(project) ?: profiles[0] + ApplicationManager.getApplication().invokeLater { + QRegionProfileDialog( + project, + profiles = profiles, + selectedProfile = selectedProfile + ).show() + } + } + } +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/clients/AmazonQStreamingClient.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/clients/AmazonQStreamingClient.kt index 43c6e85c8b1..bfba92510f2 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/clients/AmazonQStreamingClient.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/clients/AmazonQStreamingClient.kt @@ -16,20 +16,15 @@ import software.amazon.awssdk.services.codewhispererstreaming.model.ThrottlingEx import software.amazon.awssdk.services.codewhispererstreaming.model.ValidationException import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.warn -import software.aws.toolkits.jetbrains.core.awsClient -import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager -import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection import software.aws.toolkits.jetbrains.services.amazonq.RetryableOperation +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import java.time.Instant import java.util.concurrent.TimeoutException import java.util.concurrent.atomic.AtomicReference @Service(Service.Level.PROJECT) class AmazonQStreamingClient(private val project: Project) { - private fun connection() = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) - ?: error("Attempted to use connection while one does not exist") - - private fun streamingBearerClient() = connection().getConnectionSettings().awsClient() + private fun streamingBearerClient() = QRegionProfileManager.getInstance().getQClient(project) suspend fun exportResultArchive( exportId: String, diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QEndpoints.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QEndpoints.kt new file mode 100644 index 00000000000..8119def3686 --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QEndpoints.kt @@ -0,0 +1,38 @@ +// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package software.aws.toolkits.jetbrains.services.amazonq.profile +import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper +import com.fasterxml.jackson.module.kotlin.readValue +import com.intellij.openapi.util.registry.Registry +import software.aws.toolkits.core.utils.getLogger +import software.aws.toolkits.core.utils.warn + +object QEndpoints { + private val LOG = getLogger() + data class QRegionEndpoint(val region: String, val endpoint: String) + + object Q_DEFAULT_SERVICE_CONFIG { + const val REGION = "us-east-1" + const val ENDPOINT = "https://codewhisperer.us-east-1.amazonaws.com/" + } + + private fun parseEndpoints(): Map { + val rawJson = Registry.get("amazon.q.endpoints.json").asString().takeIf { it.isNotBlank() } ?: return emptyMap() + return try { + val regionList: List = jacksonObjectMapper().readValue(rawJson) + regionList.associate { it.region to it.endpoint } + } catch (e: Exception) { + LOG.warn(e) { "Failed to parse amazon.q.endpoints.json: $rawJson" } + emptyMap() + } + } + + fun listRegionEndpoints(): List = parseEndpoints().map { (region, endpoint) -> QRegionEndpoint(region, endpoint) } + + fun getQEndpointWithRegion(regionId: String): String { + val all = parseEndpoints() + return all[regionId] + ?: error("No available endpoint for region=$regionId (check amazon.q.endpoints.json or default fallback)") + } +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileResources.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileResources.kt new file mode 100644 index 00000000000..79cbdb38a27 --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileResources.kt @@ -0,0 +1,39 @@ +// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package software.aws.toolkits.jetbrains.services.amazonq.profile + +import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient +import software.aws.toolkits.core.ClientConnectionSettings +import software.aws.toolkits.jetbrains.core.AwsClientManager +import software.aws.toolkits.jetbrains.core.Resource +import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider +import java.time.Duration + +/** + * Save Amazon Q Profile Resource Cache + */ +object QProfileResources { + /** + * save available Q Profile list as cache with default duration 60 s。 + */ + val LIST_REGION_PROFILES = object : Resource.Cached>() { + override val id: String = "amazonq.allProfiles" + + override fun fetch(connectionSettings: ClientConnectionSettings<*>): List { + val mappedProfiles = QEndpoints.listRegionEndpoints().flatMap { (regionKey, _) -> + val awsRegion = AwsRegionProvider.getInstance()[regionKey] ?: return@flatMap emptyList() + val client = AwsClientManager + .getInstance() + .getClient(CodeWhispererRuntimeClient::class, connectionSettings.withRegion(awsRegion)) + + client.listAvailableProfilesPaginator {} + .profiles() + .map { p -> QRegionProfile(arn = p.arn(), profileName = p.profileName() ?: "") } + } + return mappedProfiles + } + + override fun expiry(): Duration = Duration.ofSeconds(60) + } +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileSwitchIntent.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileSwitchIntent.kt new file mode 100644 index 00000000000..c1a7b4627ab --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileSwitchIntent.kt @@ -0,0 +1,19 @@ +// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package software.aws.toolkits.jetbrains.services.amazonq.profile + +/** + * 'user' -> users change the profile through Q menu + * 'auth' -> users change the profile through webview profile selector page + * 'update' -> plugin auto select the profile on users' behalf as there is only 1 profile + * 'reload' -> on plugin restart, plugin will try to reload previous selected profile + */ +enum class QProfileSwitchIntent(val value: String) { + User("user"), + Auth("auth"), + Update("update"), + Reload("reload"), ; + + override fun toString() = value +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfile.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfile.kt new file mode 100644 index 00000000000..e8181d75228 --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfile.kt @@ -0,0 +1,24 @@ +// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package software.aws.toolkits.jetbrains.services.amazonq.profile +import software.amazon.awssdk.arns.Arn +import software.aws.toolkits.core.utils.tryOrNull + +data class QRegionProfile( + var profileName: String = "", + var arn: String = "", +) { + private val parsedArn: Arn? by lazy { + tryOrNull { + Arn.fromString(arn) + } + } + val accountId: String by lazy { + parsedArn?.accountId()?.get().orEmpty() + } + + val region: String by lazy { + parsedArn?.region()?.get().orEmpty() + } +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileDialog.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileDialog.kt new file mode 100644 index 00000000000..ead6ab5aca3 --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileDialog.kt @@ -0,0 +1,98 @@ +// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package software.aws.toolkits.jetbrains.services.amazonq.profile + +import com.intellij.icons.AllIcons +import com.intellij.openapi.project.Project +import com.intellij.openapi.ui.DialogPanel +import com.intellij.openapi.ui.DialogWrapper +import com.intellij.ui.dsl.builder.BottomGap +import com.intellij.ui.dsl.builder.bind +import com.intellij.ui.dsl.builder.panel +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.help.HelpIds +import software.aws.toolkits.resources.AmazonQBundle.message +import software.aws.toolkits.telemetry.MetricResult +import software.aws.toolkits.telemetry.Telemetry +import javax.swing.JComponent + +class QRegionProfileDialog( + private var project: Project, + private var profiles: List, + private var selectedProfile: QRegionProfile, // default +) : DialogWrapper(project) { + + private val panel: DialogPanel by lazy { + panel { + row { label(message("action.q.switchProfiles.dialog.panel.text")).bold() } + .bottomGap(BottomGap.MEDIUM) + row { text(message("action.q.switchProfiles.dialog.panel.description")) } + row { + icon(AllIcons.General.Warning) + text(message("action.q.switchProfiles.dialog.panel.warning")) + } + separator().bottomGap(BottomGap.MEDIUM) + + buttonsGroup { + profiles.forEach { profile -> + row { + radioButton("", profile) + + panel { + val regionDisplay = if (profile == selectedProfile) { + "${profile.profileName} - ${profile.region} (connected)" + } else { + "${profile.profileName} - ${profile.region}" + } + row { label(regionDisplay) } + row { + label(message("action.q.switchProfiles.dialog.account.label", profile.accountId)).applyToComponent { + font = font.deriveFont(font.size2D - 2.0f) + } + } + } + }.bottomGap(BottomGap.MEDIUM) + } + }.bind({ selectedOption }, { selectedOption = it }) + + separator().bottomGap(BottomGap.MEDIUM) + } + } + private var selectedOption: QRegionProfile = selectedProfile // user selected + + init { + title = message("action.q.switchProfiles.dialog.text") + setOKButtonText(message("general.ok")) + setCancelButtonText(message("general.cancel")) + init() + } + + override fun getHelpId(): String = HelpIds.Q_SWITCH_PROFILES_DIALOG.id + override fun createCenterPanel(): JComponent = panel + override fun doOKAction() { + panel.apply() + if (selectedOption != selectedProfile) { + QRegionProfileManager.getInstance().switchProfile(project, selectedOption, intent = QProfileSwitchIntent.User) + } + close(OK_EXIT_CODE) + } + + override fun doCancelAction() { + super.doCancelAction() + val profileManager = QRegionProfileManager.getInstance() + val conn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) as? AwsBearerTokenConnection + Telemetry.amazonq.didSelectProfile.use { span -> + span.source(QProfileSwitchIntent.User.value) + .amazonQProfileRegion(profileManager.activeProfile(project)?.region ?: "not-set") + .profileCount(profiles.size) + .ssoRegion(conn?.region) + .credentialStartUrl(conn?.startUrl) + .result(MetricResult.Cancelled) + } + + close(CANCEL_EXIT_CODE) + } +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt new file mode 100644 index 00000000000..c9014d16d20 --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt @@ -0,0 +1,232 @@ +// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package software.aws.toolkits.jetbrains.services.amazonq.profile + +import com.intellij.openapi.Disposable +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.components.BaseState +import com.intellij.openapi.components.PersistentStateComponent +import com.intellij.openapi.components.Service +import com.intellij.openapi.components.State +import com.intellij.openapi.components.Storage +import com.intellij.openapi.components.service +import com.intellij.openapi.project.Project +import com.intellij.util.concurrency.annotations.RequiresBackgroundThread +import com.intellij.util.xmlb.annotations.MapAnnotation +import com.intellij.util.xmlb.annotations.Property +import software.amazon.awssdk.core.SdkClient +import software.aws.toolkits.core.TokenConnectionSettings +import software.aws.toolkits.core.utils.debug +import software.aws.toolkits.core.utils.getLogger +import software.aws.toolkits.core.utils.tryOrNull +import software.aws.toolkits.core.utils.warn +import software.aws.toolkits.jetbrains.core.AwsClientManager +import software.aws.toolkits.jetbrains.core.AwsResourceCache +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.BearerTokenProviderListener +import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider +import software.aws.toolkits.jetbrains.utils.notifyInfo +import software.aws.toolkits.resources.AmazonQBundle.message +import software.aws.toolkits.telemetry.MetricResult +import software.aws.toolkits.telemetry.Telemetry +import java.time.Duration +import java.util.Collections +import kotlin.reflect.KClass + +@Service(Service.Level.APP) +@State(name = "qProfileStates", storages = [Storage("aws.xml")]) +class QRegionProfileManager : PersistentStateComponent, Disposable { + + // Map to store connectionId to its active profile + private val connectionIdToActiveProfile = Collections.synchronizedMap(mutableMapOf()) + private val connectionIdToProfileCount = mutableMapOf() + + init { + ApplicationManager.getApplication().messageBus.connect(this) + .subscribe( + BearerTokenProviderListener.TOPIC, + object : BearerTokenProviderListener { + override fun invalidate(providerId: String) { + connectionIdToActiveProfile.remove(providerId) + connectionIdToProfileCount.remove(providerId) + } + } + ) + } + + // should be call on project startup to validate if profile is still active + @RequiresBackgroundThread + fun validateProfile(project: Project) { + val conn = getIdcConnectionOrNull(project) + val selected = activeProfile(project) ?: return + val profiles = tryOrNull { + listRegionProfiles(project) + } + + if (profiles == null || profiles.none { it.arn == selected.arn }) { + invalidateProfile(selected.arn) + switchProfile(project, null, intent = QProfileSwitchIntent.Reload) + Telemetry.amazonq.profileState.use { span -> + span.source(QProfileSwitchIntent.Reload.value) + .amazonQProfileRegion(selected.region) + .ssoRegion(conn?.region) + .credentialStartUrl(conn?.startUrl) + .result(MetricResult.Failed) + } + } + } + + fun listRegionProfiles(project: Project): List? { + val connection = getIdcConnectionOrNull(project) ?: return null + return try { + val connectionSettings = connection.getConnectionSettings() + val mappedProfiles = AwsResourceCache.getInstance().getResourceNow( + resource = QProfileResources.LIST_REGION_PROFILES, + connectionSettings = connectionSettings, + timeout = Duration.ofSeconds(30), + useStale = true, + forceFetch = false + ) + if (mappedProfiles.size == 1) { + switchProfile(project, mappedProfiles.first(), intent = QProfileSwitchIntent.Update) + } + mappedProfiles.takeIf { it.isNotEmpty() }?.also { + connectionIdToProfileCount[connection.id] = it.size + } ?: error("You don't have access to the resource") + } catch (e: Exception) { + LOG.warn(e) { "Failed to list region profiles: ${e.message}" } + throw e + } + } + + fun activeProfile(project: Project): QRegionProfile? = getIdcConnectionOrNull(project)?.let { connectionIdToActiveProfile[it.id] } + + fun hasValidConnectionButNoActiveProfile(project: Project): Boolean = getIdcConnectionOrNull(project) != null && activeProfile(project) == null + + fun switchProfile(project: Project, newProfile: QRegionProfile?, intent: QProfileSwitchIntent) { + val conn = getIdcConnectionOrNull(project) ?: return + + val oldProfile = connectionIdToActiveProfile[conn.id] + if (oldProfile == newProfile) return + + connectionIdToActiveProfile[conn.id] = newProfile + LOG.debug { "Switch from profile $oldProfile to $newProfile for project ${project.name}" } + + if (newProfile != null) { + if (intent == QProfileSwitchIntent.User || intent == QProfileSwitchIntent.Auth) { + notifyInfo( + title = message("action.q.profile.usage.text"), + content = message("action.q.profile.usage", newProfile.profileName), + project = project + ) + + Telemetry.amazonq.didSelectProfile.use { span -> + span.source(intent.value) + .amazonQProfileRegion(newProfile.region) + .profileCount(connectionIdToProfileCount[conn.id]) + .ssoRegion(conn.region) + .credentialStartUrl(conn.startUrl) + .result(MetricResult.Succeeded) + } + } else { + Telemetry.amazonq.profileState.use { span -> + span.source(intent.value) + .amazonQProfileRegion(newProfile.region) + .ssoRegion(conn.region) + .credentialStartUrl(conn.startUrl) + .result(MetricResult.Succeeded) + } + } + } + + project.messageBus + .syncPublisher(QRegionProfileSelectedListener.TOPIC) + .onProfileSelected(project, newProfile) + } + + private fun invalidateProfile(arn: String) { + val updated = connectionIdToActiveProfile.filterValues { it.arn != arn } + connectionIdToActiveProfile.clear() + connectionIdToActiveProfile.putAll(updated) + } + + // for each idc connection, user should have a profile, otherwise should show the profile selection error page + fun isPendingProfileSelection(project: Project): Boolean = getIdcConnectionOrNull(project)?.let { conn -> + val profileCounts = connectionIdToProfileCount[conn.id] ?: 0 + val activeProfile = connectionIdToActiveProfile[conn.id] + profileCounts == 0 || (profileCounts > 1 && activeProfile?.arn.isNullOrEmpty()) + } ?: false + + fun shouldDisplayProfileInfo(project: Project): Boolean = getIdcConnectionOrNull(project)?.let { conn -> + (connectionIdToProfileCount[conn.id] ?: 0) > 1 + } ?: false + + fun getQClientSettings(project: Project): TokenConnectionSettings { + val conn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) + if (conn !is AwsBearerTokenConnection) { + error("not a bearer connection") + } + + val settings = conn.getConnectionSettings() + val awsRegion = AwsRegionProvider.getInstance()[QEndpoints.Q_DEFAULT_SERVICE_CONFIG.REGION] ?: error("unknown region from Q default service config") + + // TODO: different window should be able to select different profile + return activeProfile(project)?.let { profile -> + AwsRegionProvider.getInstance()[profile.region]?.let { region -> + settings.withRegion(region) + } + } ?: settings.withRegion(awsRegion) + } + + inline fun getQClient(project: Project): T = getQClient(project, T::class) + + fun getQClient(project: Project, sdkClass: KClass): T { + val settings = getQClientSettings(project) + val client = AwsClientManager.getInstance().getClient(sdkClass, settings) + return client + } + + private fun getIdcConnectionOrNull(project: Project): AwsBearerTokenConnection? { + val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) + if (connection is AwsBearerTokenConnection && !connection.isSono()) { + return connection + } + return null + } + + companion object { + private val LOG = getLogger() + fun getInstance(): QRegionProfileManager = service() + } + + override fun dispose() {} + + override fun getState(): QProfileState { + val state = QProfileState() + state.connectionIdToActiveProfile.putAll(this.connectionIdToActiveProfile) + state.connectionIdToProfileList.putAll(this.connectionIdToProfileCount) + return state + } + + override fun loadState(state: QProfileState) { + connectionIdToActiveProfile.clear() + connectionIdToActiveProfile.putAll(state.connectionIdToActiveProfile) + + connectionIdToProfileCount.clear() + connectionIdToProfileCount.putAll(state.connectionIdToProfileList) + } +} + +class QProfileState : BaseState() { + @get:Property + @get:MapAnnotation + val connectionIdToActiveProfile by map() + + @get:Property + @get:MapAnnotation + val connectionIdToProfileList by map() +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileSelectedListener.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileSelectedListener.kt new file mode 100644 index 00000000000..f107d883169 --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileSelectedListener.kt @@ -0,0 +1,16 @@ +// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package software.aws.toolkits.jetbrains.services.amazonq.profile + +import com.intellij.openapi.project.Project +import com.intellij.util.messages.Topic + +interface QRegionProfileSelectedListener { + companion object { + @Topic.ProjectLevel + val TOPIC = Topic.create("QRegionProfileSelected", QRegionProfileSelectedListener::class.java) + } + + fun onProfileSelected(project: Project, profile: QRegionProfile?) +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt index cb9d8b9362a..475618747c0 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt @@ -74,8 +74,8 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En ) // TODO: move to LspMessage.kt + @JsonIgnoreProperties(ignoreUnknown = true) data class Usage( - @JsonIgnoreProperties(ignoreUnknown = true) @JsonProperty("memoryUsage") val memoryUsage: Int? = null, @JsonProperty("cpuUsage") @@ -83,8 +83,8 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En ) // TODO: move to LspMessage.kt + @JsonIgnoreProperties(ignoreUnknown = true) data class Chunk( - @JsonIgnoreProperties(ignoreUnknown = true) @JsonProperty("filePath") val filePath: String? = null, @JsonProperty("content") diff --git a/plugins/amazonq/src/main/resources/META-INF/plugin.xml b/plugins/amazonq/src/main/resources/META-INF/plugin.xml index 3f359ac25ea..3d8f95579e5 100644 --- a/plugins/amazonq/src/main/resources/META-INF/plugin.xml +++ b/plugins/amazonq/src/main/resources/META-INF/plugin.xml @@ -85,7 +85,11 @@ + defaultValue="" restartRequired="true"/> + diff --git a/plugins/core/jetbrains-community/resources/telemetryOverride.json b/plugins/core/jetbrains-community/resources/telemetryOverride.json index 5ad09ab9c0e..d0f27341496 100644 --- a/plugins/core/jetbrains-community/resources/telemetryOverride.json +++ b/plugins/core/jetbrains-community/resources/telemetryOverride.json @@ -1,5 +1,20 @@ { "types": [ + { + "name": "amazonQProfileRegion", + "type": "string", + "description": "Region of the Q Profile associated with a metric\n- \"n/a\" if metric is not associated with a profile or region.\n- \"not-set\" if metric is associated with a profile, but profile is unknown." + }, + { + "name": "ssoRegion", + "type": "string", + "description": "Region of the current SSO connection. Typically associated with credentialStartUrl\n- \"n/a\" if metric is not associated with a region.\n- \"not-set\" if metric is associated with a region, but region is unknown." + }, + { + "name": "profileCount", + "type": "int", + "description": "The number of profiles that were available to choose from" + }, { "name": "amazonqIndexFileSizeInMB", "type": "int", @@ -105,6 +120,31 @@ } ], "metrics": [ + { + "name": "amazonq_didSelectProfile", + "description": "Emitted after the user's Q Profile has been set, whether the user was prompted with a dialog, or a profile was automatically assigned after signing in.", + "metadata": [ + { "type": "source" }, + { "type": "amazonQProfileRegion" }, + { "type": "result" }, + { "type": "ssoRegion", "required": false }, + { "type": "credentialStartUrl", "required": false }, + { "type": "profileCount", "required": false } + ], + "passive": true + }, + { + "name": "amazonq_profileState", + "description": "Indicates a change in the user's Q Profile state", + "metadata": [ + { "type": "source" }, + { "type": "amazonQProfileRegion" }, + { "type": "result" }, + { "type": "ssoRegion", "required": false }, + { "type": "credentialStartUrl", "required": false } + ], + "passive": true + }, { "name": "amazonq_indexWorkspace", "description": "Indexing of local workspace", diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/help/HelpIds.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/help/HelpIds.kt index 1188db0a60c..73bf084b04b 100644 --- a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/help/HelpIds.kt +++ b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/help/HelpIds.kt @@ -125,6 +125,12 @@ enum class HelpIds(shortId: String, val url: String) { "ToolkitAddConnectionsDialog", "https://docs.aws.amazon.com/toolkit-for-jetbrains/latest/userguide/setup-credentials.html" ), + + // TODO: update this + Q_SWITCH_PROFILES_DIALOG( + "QSwitchProfilesDialog", + "https://aws.amazon.com/q/developer/" + ), ; val id = "$HELP_ID_PREFIX.$shortId" diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/BrowserMessage.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/BrowserMessage.kt index 01cc00c48fa..32f0a0e1576 100644 --- a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/BrowserMessage.kt +++ b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/BrowserMessage.kt @@ -25,7 +25,9 @@ import com.fasterxml.jackson.annotation.JsonTypeInfo JsonSubTypes.Type(value = BrowserMessage.CancelLogin::class, name = "cancelLogin"), JsonSubTypes.Type(value = BrowserMessage.Signout::class, name = "signout"), JsonSubTypes.Type(value = BrowserMessage.Reauth::class, name = "reauth"), - JsonSubTypes.Type(value = BrowserMessage.SendUiClickTelemetry::class, name = "sendUiClickTelemetry") + JsonSubTypes.Type(value = BrowserMessage.SendUiClickTelemetry::class, name = "sendUiClickTelemetry"), + JsonSubTypes.Type(value = BrowserMessage.SwitchProfile::class, name = "switchProfile"), + JsonSubTypes.Type(value = BrowserMessage.PublishWebviewTelemetry::class, name = "webviewTelemetry") ) sealed interface BrowserMessage { @@ -57,5 +59,14 @@ sealed interface BrowserMessage { object Reauth : BrowserMessage + data class SwitchProfile( + val profileName: String, + val accountId: String, + val region: String, + val arn: String, + ) : BrowserMessage + data class SendUiClickTelemetry(val signInOptionClicked: String?) : BrowserMessage + + data class PublishWebviewTelemetry(val event: String) : BrowserMessage } diff --git a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/LoginBrowser.kt b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/LoginBrowser.kt index 86970a2dfd2..2c9049f5b72 100644 --- a/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/LoginBrowser.kt +++ b/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/webview/LoginBrowser.kt @@ -406,6 +406,26 @@ abstract class LoginBrowser( return false } + // TODO: should test via handleMessage, however because we can't initiate Q/ToolkitLoginBrowser in test due to jcef not supported in test env + // plus handleMessage is abstract so as a interim, exposing it for testing purpose + @VisibleForTesting + fun publishTelemetry(message: BrowserMessage.PublishWebviewTelemetry) { + val jsonNode = this.objectMapper.readTree(message.event) ?: return + if (jsonNode["metricName"].asText() == "toolkit_didLoadModule") { + val moduleNode = jsonNode["module"] ?: return + val resultNode = jsonNode["result"] ?: return + val result = MetricResult.from(resultNode.asText()) + val reasonNode = jsonNode["reason"] + val durationNode = jsonNode["duration"] + Telemetry.toolkit.didLoadModule.use { span -> + span.module(moduleNode.asText()) + span.result(result) + span.reason(reasonNode?.asText()) + span.duration(durationNode?.asDouble()) + } + } + } + companion object { private val LOG = getLogger() fun getWebviewHTML(webScriptUri: String, query: JBCefJSQuery): String { diff --git a/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/BrowserMessageTest.kt b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/BrowserMessageTest.kt index 0ee085f52e6..883945535d4 100644 --- a/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/BrowserMessageTest.kt +++ b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/BrowserMessageTest.kt @@ -161,6 +161,19 @@ class BrowserMessageTest { signInOptionClicked = null ) ) + + assertDeserializedInstanceOf( + """ + { + "command": "webviewTelemetry", + "event": "{ \"metricName\": \"foo\" }" + } + """.trimIndent() + ).isEqualTo( + BrowserMessage.PublishWebviewTelemetry( + event = "{ \"metricName\": \"foo\" }" + ) + ) } @Test diff --git a/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/LoginBrowserTest.kt b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/LoginBrowserTest.kt new file mode 100644 index 00000000000..f5ed7e0f2c5 --- /dev/null +++ b/plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/LoginBrowserTest.kt @@ -0,0 +1,159 @@ +// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package software.aws.toolkits.jetbrains.core + +import com.intellij.openapi.project.Project +import com.intellij.testFramework.ProjectExtension +import com.intellij.ui.jcef.JBCefBrowserBase +import com.intellij.ui.jcef.JBCefJSQuery +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.BeforeEach +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.extension.RegisterExtension +import org.mockito.kotlin.argumentCaptor +import org.mockito.kotlin.mock +import org.mockito.kotlin.times +import org.mockito.kotlin.verify +import software.aws.toolkits.core.telemetry.MetricEvent +import software.aws.toolkits.jetbrains.core.webview.BrowserMessage +import software.aws.toolkits.jetbrains.core.webview.BrowserState +import software.aws.toolkits.jetbrains.core.webview.LoginBrowser +import software.aws.toolkits.jetbrains.services.telemetry.MockTelemetryServiceExtension +import kotlin.test.assertNotNull + +class TestLoginBrowser(project: Project) : LoginBrowser(project, "", "") { + // test env can't initiate a real jcef and will throw error + override val jcefBrowser: JBCefBrowserBase + get() = mock() + + override fun handleBrowserMessage(message: BrowserMessage?) {} + + override fun prepareBrowser(state: BrowserState) {} + + override fun loadWebView(query: JBCefJSQuery) {} +} + +class LoginBrowserTest { + private lateinit var sut: TestLoginBrowser + private val project: Project + get() = projectExtension.project + + @JvmField + @RegisterExtension + val mockTelemetryService = MockTelemetryServiceExtension() + + companion object { + @JvmField + @RegisterExtension + val projectExtension = ProjectExtension() + } + + @BeforeEach + fun setup() { + sut = TestLoginBrowser(project) + } + + @Test + fun `publish telemetry happy path`() { + val load = """ + { + "metricName": "toolkit_didLoadModule", + "module": "login", + "result": "Succeeded", + "duration": "0" + } + """.trimIndent() + val message = BrowserMessage.PublishWebviewTelemetry(load) + sut.publishTelemetry(message) + + mockTelemetryService.batcher() + argumentCaptor { + verify(mockTelemetryService.batcher()).enqueue(capture()) + val event = firstValue.data.find { it.name == "toolkit_didLoadModule" } + assertNotNull(event) + assertThat(event) + .matches { it.metadata["module"] == "login" } + .matches { it.metadata["result"] == "Succeeded" } + .matches { it.metadata["duration"] == "0.0" } + } + } + + @Test + fun `publish telemetry error path`() { + val load = """ + { + "metricName": "toolkit_didLoadModule", + "module": "login", + "result": "Failed", + "reason": "unexpected error" + } + """.trimIndent() + val message = BrowserMessage.PublishWebviewTelemetry(load) + sut.publishTelemetry(message) + + mockTelemetryService.batcher() + argumentCaptor { + verify(mockTelemetryService.batcher()).enqueue(capture()) + val event = firstValue.data.find { it.name == "toolkit_didLoadModule" } + assertNotNull(event) + assertThat(event) + .matches { it.metadata["module"] == "login" } + .matches { it.metadata["result"] == "Failed" } + .matches { it.metadata["reason"] == "unexpected error" } + } + } + + @Test + fun `missing required field will do nothing`() { + val load = """ + { + "metricName": "toolkit_didLoadModule" + } + """.trimIndent() + val message = BrowserMessage.PublishWebviewTelemetry(load) + sut.publishTelemetry(message) + + val load1 = """ + { + "metricName": "toolkit_didLoadModule", + "module": "login" + } + """.trimIndent() + val message1 = BrowserMessage.PublishWebviewTelemetry(load1) + sut.publishTelemetry(message1) + + val load2 = """ + { + "metricName": "toolkit_didLoadModule", + "result": "Failed" + } + """.trimIndent() + val message2 = BrowserMessage.PublishWebviewTelemetry(load2) + sut.publishTelemetry(message2) + + mockTelemetryService.batcher() + argumentCaptor { + verify(mockTelemetryService.batcher(), times(0)).enqueue(capture()) + } + } + + @Test + fun `metricName doesn't match will do nothing`() { + val load = """ + { + "metricName": "foo", + "module": "login", + "result": "Failed", + "reason": "unexpected error" + } + """.trimIndent() + val message = BrowserMessage.PublishWebviewTelemetry(load) + sut.publishTelemetry(message) + + mockTelemetryService.batcher() + argumentCaptor { + verify(mockTelemetryService.batcher(), times(0)).enqueue(capture()) + } + } +} diff --git a/plugins/core/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/core/MockResourceCache.kt b/plugins/core/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/core/MockResourceCache.kt index 66d5249ce53..78cfbe4409c 100644 --- a/plugins/core/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/core/MockResourceCache.kt +++ b/plugins/core/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/core/MockResourceCache.kt @@ -13,7 +13,6 @@ import org.junit.jupiter.api.extension.BeforeEachCallback import org.junit.jupiter.api.extension.ExtensionContext import org.junit.runner.Description import software.aws.toolkits.core.ClientConnectionSettings -import software.aws.toolkits.core.ConnectionSettings import software.aws.toolkits.core.credentials.ToolkitAuthenticationProvider import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider import software.aws.toolkits.core.credentials.ToolkitCredentialsProvider @@ -182,8 +181,12 @@ interface MockResourceCacheInterface { addEntry(project, resourceId, CompletableFuture.failedFuture(throws)) } - fun addEntry(connectionSettings: ConnectionSettings, resource: Resource.Cached, value: CompletableFuture) { - addEntry(resource, connectionSettings.region.id, connectionSettings.credentials.id, value) + fun addEntry(connectionSettings: ClientConnectionSettings<*>, resource: Resource.Cached, value: T) { + addEntry(resource, connectionSettings.region.id, connectionSettings.providerId, value) + } + + fun addEntry(connectionSettings: ClientConnectionSettings<*>, resource: Resource.Cached, value: CompletableFuture) { + addEntry(resource, connectionSettings.region.id, connectionSettings.providerId, value) } fun addEntry(resource: Resource.Cached, regionId: String, credentialsId: String, value: T) { diff --git a/plugins/core/sdk-codegen/codegen-resources/codewhispererruntime/paginators-1.json b/plugins/core/sdk-codegen/codegen-resources/codewhispererruntime/paginators-1.json index c0860a2f5a5..381687933c1 100644 --- a/plugins/core/sdk-codegen/codegen-resources/codewhispererruntime/paginators-1.json +++ b/plugins/core/sdk-codegen/codegen-resources/codewhispererruntime/paginators-1.json @@ -11,9 +11,21 @@ "limit_key": "maxResults", "result_key": "customizations" }, + "ListAvailableProfiles": { + "input_token": "nextToken", + "output_token": "nextToken", + "limit_key": "maxResults", + "result_key": "profiles" + }, "ListCodeAnalysisFindings": { "input_token": "nextToken", "output_token": "nextToken" + }, + "ListWorkspaceMetadata": { + "input_token": "nextToken", + "output_token": "nextToken", + "limit_key": "maxResults", + "result_key": "workspaces" } } } diff --git a/plugins/core/webview/src/ideClient.ts b/plugins/core/webview/src/ideClient.ts index 0abed7d8527..902e854d33f 100644 --- a/plugins/core/webview/src/ideClient.ts +++ b/plugins/core/webview/src/ideClient.ts @@ -2,7 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 import {Store} from "vuex"; -import {IdcInfo, Region, Stage, State, BrowserSetupData, AwsBearerTokenConnection} from "./model"; +import {IdcInfo, Region, Stage, State, BrowserSetupData, AwsBearerTokenConnection, Profile} from "./model"; +import {WebviewTelemetry} from './webviewTelemetry' export class IdeClient { constructor(private readonly store: Store) {} @@ -10,13 +11,18 @@ export class IdeClient { // TODO: design and improve the API here prepareUi(state: BrowserSetupData) { + WebviewTelemetry.instance.reset() console.log('browser is preparing UI with state ', state) this.store.commit('setStage', state.stage) + // hack as window.onerror don't have access to vuex store + void ((window as any).uiState = state.stage) + WebviewTelemetry.instance.willShowPage(state.stage) this.store.commit('setSsoRegions', state.regions) this.updateLastLoginIdcInfo(state.idcInfo) this.store.commit("setCancellable", state.cancellable) this.store.commit("setFeature", state.feature) - + this.store.commit('setProfiles', state.profiles); + this.store.commit("setErrorMessage", state.errorMessage) const existConnections = state.existConnections.map(it => { return { sessionName: it.sessionName, @@ -31,6 +37,13 @@ export class IdeClient { this.updateAuthorization(undefined) } + handleProfiles(profilesData: { profiles: Profile[] }) { + this.store.commit('setStage', 'PROFILE_SELECT') + console.debug("received profile data") + const availableProfiles: Profile[] = profilesData.profiles; + this.store.commit('setProfiles', availableProfiles); + } + updateAuthorization(code: string | undefined) { this.store.commit('setAuthorizationCode', code) // TODO: mutage stage to AUTHing here probably makes life easier diff --git a/plugins/core/webview/src/model.ts b/plugins/core/webview/src/model.ts index 424628e011b..d30f7f85eb7 100644 --- a/plugins/core/webview/src/model.ts +++ b/plugins/core/webview/src/model.ts @@ -7,7 +7,9 @@ export type BrowserSetupData = { idcInfo: IdcInfo, cancellable: boolean, feature: string, - existConnections: AwsBearerTokenConnection[] + existConnections: AwsBearerTokenConnection[], + profiles: Profile[], + errorMessage: string } // plugin interface [AwsBearerTokenConnection] @@ -26,7 +28,8 @@ export type Stage = 'CONNECTED' | 'AUTHENTICATING' | 'AWS_PROFILE' | - 'REAUTH' + 'REAUTH' | + 'PROFILE_SELECT' export type Feature = 'Q' | 'codecatalyst' | 'awsExplorer' @@ -50,7 +53,10 @@ export interface State { lastLoginIdcInfo: IdcInfo, feature: Feature, cancellable: boolean, - existingConnections: AwsBearerTokenConnection[] + existingConnections: AwsBearerTokenConnection[], + profiles: Profile[], + selectedProfile: Profile | undefined, + errorMessage: string | undefined } export enum LoginIdentifier { @@ -67,6 +73,15 @@ export interface LoginOption { requiresBrowser(): boolean } +export interface Profile { + profileName: string + accountId: string + region: string + arn: String +} + +export const GENERIC_PROFILE_LOAD_ERROR = "We couldn't load your Q Developer profiles. Please try again."; + export class LongLivedIAM implements LoginOption { id: LoginIdentifier = LoginIdentifier.IAM_CREDENTIAL diff --git a/plugins/core/webview/src/q-ui/components/login.vue b/plugins/core/webview/src/q-ui/components/login.vue index 260c44e6e95..6b31e1efe4d 100644 --- a/plugins/core/webview/src/q-ui/components/login.vue +++ b/plugins/core/webview/src/q-ui/components/login.vue @@ -4,8 +4,7 @@