Skip to content

Commit 6cfc158

Browse files
samgst-amazonleigaol
authored andcommitted
implement token change listener (aws#5423)
1 parent 7973775 commit 6cfc158

File tree

3 files changed

+47
-2
lines changed

3 files changed

+47
-2
lines changed

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ import software.aws.toolkits.core.utils.getLogger
4545
import software.aws.toolkits.core.utils.info
4646
import software.aws.toolkits.core.utils.warn
4747
import software.aws.toolkits.jetbrains.isDeveloperMode
48+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.auth.DefaultAuthCredentialsService
4849
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
4950
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.createExtendedClientMetadata
5051
import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata
@@ -303,6 +304,8 @@ private class AmazonQServerInstance(private val project: Project, private val cs
303304
}
304305
languageServer.initialized(InitializedParams())
305306
}
307+
308+
DefaultAuthCredentialsService(project, encryptionManager, this)
306309
}
307310

308311
override fun dispose() {

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

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,14 @@
33

44
package software.aws.toolkits.jetbrains.services.amazonq.lsp.auth
55

6+
import com.intellij.openapi.Disposable
67
import com.intellij.openapi.project.Project
78
import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
9+
import software.aws.toolkits.core.TokenConnectionSettings
10+
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
11+
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
12+
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
13+
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
814
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
915
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
1016
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.BearerCredentials
@@ -15,7 +21,31 @@ import java.util.concurrent.CompletableFuture
1521
class DefaultAuthCredentialsService(
1622
private val project: Project,
1723
private val encryptionManager: JwtEncryptionManager,
18-
) : AuthCredentialsService {
24+
serverInstance: Disposable,
25+
) : AuthCredentialsService,
26+
BearerTokenProviderListener {
27+
init {
28+
project.messageBus.connect(serverInstance).subscribe(BearerTokenProviderListener.TOPIC, this)
29+
}
30+
31+
override fun onChange(providerId: String, newScopes: List<String>?) {
32+
val connection = ToolkitConnectionManager.getInstance(project)
33+
.activeConnectionForFeature(QConnection.getInstance())
34+
?: return
35+
36+
val provider = (connection.getConnectionSettings() as? TokenConnectionSettings)
37+
?.tokenProvider
38+
?.delegate as? BearerTokenProvider
39+
?: return
40+
41+
provider.currentToken()?.accessToken?.let { token ->
42+
updateTokenCredentials(token, false)
43+
}
44+
}
45+
46+
override fun invalidate(providerId: String) {
47+
deleteTokenCredentials()
48+
}
1949

2050
override fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture<ResponseMessage> {
2151
val token = if (encrypted) {

plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.kt

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,15 @@
33

44
package software.aws.toolkits.jetbrains.services.amazonq.lsp.auth
55

6+
import com.intellij.openapi.Disposable
67
import com.intellij.openapi.components.serviceIfCreated
78
import com.intellij.openapi.project.Project
9+
import com.intellij.util.messages.MessageBus
10+
import com.intellij.util.messages.MessageBusConnection
811
import io.mockk.every
12+
import io.mockk.just
913
import io.mockk.mockk
14+
import io.mockk.runs
1015
import io.mockk.verify
1116
import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
1217
import org.junit.Before
@@ -42,7 +47,14 @@ class DefaultAuthCredentialsServiceTest {
4247
func.invoke(mockLanguageServer)
4348
}
4449

45-
sut = DefaultAuthCredentialsService(project, this.mockEncryptionManager)
50+
// Mock message bus
51+
val messageBus = mockk<MessageBus>()
52+
every { project.messageBus } returns messageBus
53+
val mockConnection = mockk<MessageBusConnection>()
54+
every { messageBus.connect(any<Disposable>()) } returns mockConnection
55+
every { mockConnection.subscribe(any(), any()) } just runs
56+
57+
sut = DefaultAuthCredentialsService(project, this.mockEncryptionManager, mockk())
4658
}
4759

4860
@Test

0 commit comments

Comments
 (0)