From 7125dbe2f5f60ae102936e2f45f16965181f0ed6 Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Thu, 17 Apr 2025 13:48:21 -0700 Subject: [PATCH 1/3] updateConfiguration --- .../amazonq/lsp/AmazonQLanguageServer.kt | 4 +++ .../lsp/auth/DefaultAuthCredentialsService.kt | 26 ++++++++++++++++++- .../lsp/model/aws/LspServerConfigurations.kt | 7 +++++ 3 files changed, 36 insertions(+), 1 deletion(-) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt index 18cf95d3973..2396e273f18 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt @@ -9,6 +9,7 @@ import org.eclipse.lsp4j.jsonrpc.services.JsonRequest import org.eclipse.lsp4j.services.LanguageServer 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.lsp.model.aws.UpdateConfigurationParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.dependencies.DidChangeDependencyPathsParams import java.util.concurrent.CompletableFuture @@ -29,4 +30,7 @@ interface AmazonQLanguageServer : LanguageServer { @JsonRequest("aws/getConfigurationFromServer") fun getConfigurationFromServer(params: GetConfigurationFromServerParams): CompletableFuture + + @JsonRequest("aws/updateConfiguration") + fun updateConfiguration(params: UpdateConfigurationParams): CompletableFuture } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt index c5be0abe56c..1605906b2a5 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt @@ -15,9 +15,14 @@ import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenPr import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerConfigurations +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.UpdateConfigurationParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.BearerCredentials import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayloadData +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 java.util.concurrent.CompletableFuture @@ -28,16 +33,19 @@ class DefaultAuthCredentialsService( serverInstance: Disposable, ) : AuthCredentialsService, BearerTokenProviderListener, - ToolkitConnectionManagerListener { + ToolkitConnectionManagerListener, + QRegionProfileSelectedListener { init { project.messageBus.connect(serverInstance).apply { subscribe(BearerTokenProviderListener.TOPIC, this@DefaultAuthCredentialsService) subscribe(ToolkitConnectionManagerListener.TOPIC, this@DefaultAuthCredentialsService) + subscribe(QRegionProfileSelectedListener.TOPIC, this@DefaultAuthCredentialsService) } if (isQConnected(project) && !isQExpired(project)) { updateTokenFromActiveConnection() + updateConfiguration() } } @@ -108,4 +116,20 @@ class DefaultAuthCredentialsService( encrypted = false ) } + + override fun onProfileSelected(project: Project, profile: QRegionProfile?) { + updateConfiguration() + } + + private fun updateConfiguration(): CompletableFuture { + val payload = UpdateConfigurationParams( + section = "aws.q", + settings = mapOf( + "profileArn" to QRegionProfileManager.getInstance().activeProfile(project)?.arn + ) + ) + return AmazonQLspService.executeIfRunning(project) { server -> + server.updateConfiguration(payload) + } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running"))) + } } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/LspServerConfigurations.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/LspServerConfigurations.kt index e762b938d28..a0f23875b62 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/LspServerConfigurations.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/LspServerConfigurations.kt @@ -8,3 +8,10 @@ data class WorkspaceInfo(val workspaceRoot: String, val workspaceId: String) // This represents the entire array data class LspServerConfigurations(val workspaces: List) + +data class UpdateConfigurationParams( + val section: String, + val settings: LSPAny, +) + +typealias LSPAny = Any? From 8970dcfef9c98f655c419a4e81a8a74be17bf5a6 Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Thu, 17 Apr 2025 14:37:01 -0700 Subject: [PATCH 2/3] fix race condition --- .../lsp/auth/DefaultAuthCredentialsService.kt | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt index 1605906b2a5..6c3ddef1986 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt @@ -45,7 +45,9 @@ class DefaultAuthCredentialsService( if (isQConnected(project) && !isQExpired(project)) { updateTokenFromActiveConnection() - updateConfiguration() + .thenRun { + updateConfiguration() + } } } @@ -78,22 +80,23 @@ class DefaultAuthCredentialsService( updateTokenFromConnection(newConnection) } - private fun updateTokenFromActiveConnection() { + private fun updateTokenFromActiveConnection(): CompletableFuture { val connection = ToolkitConnectionManager.getInstance(project) .activeConnectionForFeature(QConnection.getInstance()) - ?: return + ?: return CompletableFuture.failedFuture(IllegalStateException("No active Q connection")) - updateTokenFromConnection(connection) + return updateTokenFromConnection(connection) } - private fun updateTokenFromConnection(connection: ToolkitConnection) { - (connection.getConnectionSettings() as? TokenConnectionSettings) + private fun updateTokenFromConnection(connection: ToolkitConnection): CompletableFuture { + return (connection.getConnectionSettings() as? TokenConnectionSettings) ?.tokenProvider ?.delegate ?.let { it as? BearerTokenProvider } ?.currentToken() ?.accessToken ?.let { token -> updateTokenCredentials(token, true) } + ?:CompletableFuture.failedFuture(IllegalStateException("Unable to get token from connection")) } override fun invalidate(providerId: String) { From a67f1c04d4884528e73060808c9412467cceca45 Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Thu, 17 Apr 2025 14:40:49 -0700 Subject: [PATCH 3/3] detekt --- .../amazonq/lsp/auth/DefaultAuthCredentialsService.kt | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt index 6c3ddef1986..d3a99a1f4fe 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt @@ -88,16 +88,15 @@ class DefaultAuthCredentialsService( return updateTokenFromConnection(connection) } - private fun updateTokenFromConnection(connection: ToolkitConnection): CompletableFuture { - return (connection.getConnectionSettings() as? TokenConnectionSettings) + private fun updateTokenFromConnection(connection: ToolkitConnection): CompletableFuture = + (connection.getConnectionSettings() as? TokenConnectionSettings) ?.tokenProvider ?.delegate ?.let { it as? BearerTokenProvider } ?.currentToken() ?.accessToken ?.let { token -> updateTokenCredentials(token, true) } - ?:CompletableFuture.failedFuture(IllegalStateException("Unable to get token from connection")) - } + ?: CompletableFuture.failedFuture(IllegalStateException("Unable to get token from connection")) override fun invalidate(providerId: String) { deleteTokenCredentials()