deps(lsp): implement aws/updateConfiguration (#5600)

Author: samgst-amazon

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<LspServerConfigurations>
+
+    @JsonRequest("aws/updateConfiguration")
+    fun updateConfiguration(params: UpdateConfigurationParams): CompletableFuture<LspServerConfigurations>
 }

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,23 +80,23 @@ class DefaultAuthCredentialsService(
         updateTokenFromConnection(newConnection)
     }
 
-    private fun updateTokenFromActiveConnection() {
+    private fun updateTokenFromActiveConnection(): CompletableFuture<ResponseMessage> {
         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<ResponseMessage> =
         (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) {
         deleteTokenCredentials()
@@ -108,4 +118,20 @@ class DefaultAuthCredentialsService(
                 encrypted = false
             )
         }
+
+    override fun onProfileSelected(project: Project, profile: QRegionProfile?) {
+        updateConfiguration()
+    }
+
+    private fun updateConfiguration(): CompletableFuture<LspServerConfigurations> {
+        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")))
+    }
 }

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<WorkspaceInfo>)
+
+data class UpdateConfigurationParams(
+    val section: String,
+    val settings: LSPAny,
+)
+
+typealias LSPAny = Any?