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..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 @@ -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,21 @@ 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() + .thenRun { + updateConfiguration() + } } } @@ -70,15 +80,15 @@ 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) { + private fun updateTokenFromConnection(connection: ToolkitConnection): CompletableFuture = (connection.getConnectionSettings() as? TokenConnectionSettings) ?.tokenProvider ?.delegate @@ -86,7 +96,7 @@ class DefaultAuthCredentialsService( ?.currentToken() ?.accessToken ?.let { token -> updateTokenCredentials(token, true) } - } + ?: CompletableFuture.failedFuture(IllegalStateException("Unable to get token from connection")) override fun invalidate(providerId: String) { deleteTokenCredentials() @@ -108,4 +118,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?