feat(amazonq): Implement aws/credentials/token messages (#5410)

* implement aws/credentials/token messages * detekt * feedback
aws · Feb 26, 2025 · 35b0424 · 35b0424
1 parent 179aea2
commit 35b0424
Show file tree

Hide file tree

Showing 5 changed files with 164 additions and 6 deletions.
diff --git a/...mmunity/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt b/...mmunity/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt
@@ -4,6 +4,7 @@
 package software.aws.toolkits.jetbrains.services.amazonq.lsp
 
 import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
+import org.eclipse.lsp4j.jsonrpc.services.JsonNotification
 import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
 import org.eclipse.lsp4j.services.LanguageServer
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
@@ -16,4 +17,7 @@ import java.util.concurrent.CompletableFuture
 interface AmazonQLanguageServer : LanguageServer {
     @JsonRequest("aws/credentials/token/update")
     fun updateTokenCredentials(payload: UpdateCredentialsPayload): CompletableFuture<ResponseMessage>
+
+    @JsonNotification("aws/credentials/token/delete")
+    fun deleteTokenCredentials(): CompletableFuture<Unit>
 }
diff --git a/...s-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt b/...s-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt
@@ -154,28 +154,26 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
         instance = start()
     }
 
-    suspend fun execute(runnable: suspend (AmazonQLanguageServer) -> Unit) {
+    suspend fun<T> execute(runnable: suspend (AmazonQLanguageServer) -> T): T {
         val lsp = withTimeout(10.seconds) {
             val holder = mutex.withLock { instance }.await()
             holder.initializer.join()
 
             holder.languageServer
         }
-
-        runnable(lsp)
+        return runnable(lsp)
     }
 
-    fun executeSync(runnable: suspend (AmazonQLanguageServer) -> Unit) {
+    fun<T> executeSync(runnable: suspend (AmazonQLanguageServer) -> T): T =
         runBlocking(cs.coroutineContext) {
             execute(runnable)
         }
-    }
 
     companion object {
         private val LOG = getLogger<AmazonQLspService>()
         fun getInstance(project: Project) = project.service<AmazonQLspService>()
 
-        fun executeIfRunning(project: Project, runnable: (AmazonQLanguageServer) -> Unit) =
+        fun <T> executeIfRunning(project: Project, runnable: (AmazonQLanguageServer) -> T): T? =
             project.serviceIfCreated<AmazonQLspService>()?.executeSync(runnable)
     }
 }

diff --git a/...y/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/AuthCredentialsService.kt b/...y/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/AuthCredentialsService.kt
@@ -0,0 +1,12 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.auth
+
+import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
+import java.util.concurrent.CompletableFuture
+
+interface AuthCredentialsService {
+    fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture<ResponseMessage>
+    fun deleteTokenCredentials(): CompletableFuture<Unit>
+}
diff --git a/...oftware/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt b/...oftware/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt
@@ -0,0 +1,51 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.auth
+
+import com.intellij.openapi.project.Project
+import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
+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
+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 java.util.concurrent.CompletableFuture
+
+class DefaultAuthCredentialsService(
+    private val project: Project,
+    private val encryptionManager: JwtEncryptionManager,
+) : AuthCredentialsService {
+
+    override fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture<ResponseMessage> {
+        val token = if (encrypted) {
+            encryptionManager.decrypt(accessToken)
+        } else {
+            accessToken
+        }
+
+        val payload = createUpdateCredentialsPayload(token)
+
+        return AmazonQLspService.executeIfRunning(project) { server ->
+            server.updateTokenCredentials(payload)
+        } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
+    }
+
+    override fun deleteTokenCredentials(): CompletableFuture<Unit> =
+        CompletableFuture<Unit>().also { completableFuture ->
+            AmazonQLspService.executeIfRunning(project) { server ->
+                server.deleteTokenCredentials()
+                completableFuture.complete(null)
+            } ?: completableFuture.completeExceptionally(IllegalStateException("LSP Server not running"))
+        }
+
+    private fun createUpdateCredentialsPayload(token: String): UpdateCredentialsPayload =
+        UpdateCredentialsPayload(
+            data = encryptionManager.encrypt(
+                UpdateCredentialsPayloadData(
+                    BearerCredentials(token)
+                )
+            ),
+            encrypted = true
+        )
+}
diff --git a/...are/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.kt b/...are/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.kt
@@ -0,0 +1,93 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.auth
+
+import com.intellij.openapi.components.serviceIfCreated
+import com.intellij.openapi.project.Project
+import io.mockk.every
+import io.mockk.mockk
+import io.mockk.verify
+import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
+import org.junit.Before
+import org.junit.Test
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLanguageServer
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
+import java.util.concurrent.CompletableFuture
+
+class DefaultAuthCredentialsServiceTest {
+    private lateinit var project: Project
+    private lateinit var mockLanguageServer: AmazonQLanguageServer
+    private lateinit var mockEncryptionManager: JwtEncryptionManager
+    private lateinit var sut: DefaultAuthCredentialsService
+
+    @Before
+    fun setUp() {
+        project = mockk<Project>()
+        mockLanguageServer = mockk<AmazonQLanguageServer>()
+        mockEncryptionManager = mockk<JwtEncryptionManager>()
+        every { mockEncryptionManager.encrypt(any()) } returns "mock-encrypted-data"
+
+        // Mock the service methods on Project
+        val mockLspService = mockk<AmazonQLspService>()
+        every { project.getService(AmazonQLspService::class.java) } returns mockLspService
+        every { project.serviceIfCreated<AmazonQLspService>() } returns mockLspService
+
+        // Mock the LSP service's executeSync method as a suspend function
+        every {
+            mockLspService.executeSync<CompletableFuture<ResponseMessage>>(any())
+        } coAnswers {
+            val func = firstArg<suspend (AmazonQLanguageServer) -> CompletableFuture<ResponseMessage>>()
+            func.invoke(mockLanguageServer)
+        }
+
+        sut = DefaultAuthCredentialsService(project, this.mockEncryptionManager)
+    }
+
+    @Test
+    fun `test updateTokenCredentials unencrypted success`() {
+        val token = "unencryptedToken"
+        val isEncrypted = false
+
+        every {
+            mockLanguageServer.updateTokenCredentials(any())
+        } returns CompletableFuture.completedFuture(ResponseMessage())
+
+        sut.updateTokenCredentials(token, isEncrypted)
+
+        verify(exactly = 0) {
+            mockEncryptionManager.decrypt(any())
+        }
+        verify(exactly = 1) {
+            mockLanguageServer.updateTokenCredentials(any())
+        }
+    }
+
+    @Test
+    fun `test updateTokenCredentials encrypted success`() {
+        val encryptedToken = "encryptedToken"
+        val decryptedToken = "decryptedToken"
+        val isEncrypted = true
+
+        every { mockEncryptionManager.decrypt(encryptedToken) } returns decryptedToken
+        every { mockEncryptionManager.encrypt(any()) } returns "mock-encrypted-data"
+        every {
+            mockLanguageServer.updateTokenCredentials(any())
+        } returns CompletableFuture.completedFuture(ResponseMessage())
+
+        sut.updateTokenCredentials(encryptedToken, isEncrypted)
+
+        verify(exactly = 1) { mockEncryptionManager.decrypt(encryptedToken) }
+        verify(exactly = 1) { mockLanguageServer.updateTokenCredentials(any()) }
+    }
+
+    @Test
+    fun `test deleteTokenCredentials success`() {
+        every { mockLanguageServer.deleteTokenCredentials() } returns CompletableFuture.completedFuture(Unit)
+
+        sut.deleteTokenCredentials()
+
+        verify(exactly = 1) { mockLanguageServer.deleteTokenCredentials() }
+    }
+}