emit on login

Author: samgst-amazon

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt

@@ -7,10 +7,12 @@ import com.intellij.openapi.Disposable
 import com.intellij.openapi.project.Project
 import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
 import software.aws.toolkits.core.TokenConnectionSettings
+import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
+import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
 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.credentials.BearerCredentials
@@ -23,25 +25,50 @@ class DefaultAuthCredentialsService(
     private val encryptionManager: JwtEncryptionManager,
     serverInstance: Disposable,
 ) : AuthCredentialsService,
-    BearerTokenProviderListener {
-    init {
+    BearerTokenProviderListener,
+    ToolkitConnectionManagerListener {
+
+        init {
         project.messageBus.connect(serverInstance).subscribe(BearerTokenProviderListener.TOPIC, this)
+        project.messageBus.connect(serverInstance).subscribe(ToolkitConnectionManagerListener.TOPIC, this)
+
+        val connection = ToolkitConnectionManager.getInstance(project)
+            .activeConnectionForFeature(QConnection.getInstance())
+
+        val provider = (connection?.getConnectionSettings() as? TokenConnectionSettings)
+            ?.tokenProvider
+            ?.delegate as? BearerTokenProvider
 
-        onChange("init", null)
+        provider?.currentToken()?.accessToken?.let { token ->
+            updateTokenCredentials(token, true)
+        }
     }
 
     override fun onChange(providerId: String, newScopes: List<String>?) {
-        val connection = ToolkitConnectionManager.getInstance(project)
-            .activeConnectionForFeature(QConnection.getInstance())
+        val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
             ?: return
-
         val provider = (connection.getConnectionSettings() as? TokenConnectionSettings)
             ?.tokenProvider
             ?.delegate as? BearerTokenProvider
             ?: return
 
         provider.currentToken()?.accessToken?.let { token ->
-            // assume encryption is always on
+            updateTokenCredentials(token, true)
+        }
+    }
+
+    override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
+        val qConnection = ToolkitConnectionManager.getInstance(project)
+            .activeConnectionForFeature(QConnection.getInstance())
+            ?: return
+        if (newConnection?.id != qConnection.id) return
+        // Handle new connection
+        val provider = (newConnection.getConnectionSettings() as? TokenConnectionSettings)
+            ?.tokenProvider
+            ?.delegate as? BearerTokenProvider
+            ?: return
+
+        provider.currentToken()?.accessToken?.let { token ->
             updateTokenCredentials(token, true)
         }
     }
@@ -82,4 +109,5 @@ class DefaultAuthCredentialsService(
                 encrypted = false
             )
         }
+
 }