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

plugins/amazonq/shared/jetbrains-community/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()
 }

plugins/amazonq/shared/jetbrains-community/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>
+}

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

@@ -0,0 +1,61 @@
+// 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 CompletableFuture<ResponseMessage>().also { completableFuture ->
+            AmazonQLspService.executeIfRunning(project) { server ->
+                server.updateTokenCredentials(payload)
+                    .whenComplete { response, throwable ->
+                        if (throwable != null) {
+                            completableFuture.completeExceptionally(throwable)
+                        } else {
+                            completableFuture.complete(response)
+                        }
+                    }
+            } ?: completableFuture.completeExceptionally(IllegalStateException("LSP Server not running"))
+        }
+    }
+
+    override fun deleteTokenCredentials(): CompletableFuture<Unit> {
+        return 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
+        )
+}

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

@@ -0,0 +1,96 @@
+// 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.coEvery
+import io.mockk.every
+import io.mockk.just
+import io.mockk.mockk
+import io.mockk.runs
+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
+        coEvery {
+            mockLspService.executeSync(any())
+        } coAnswers {
+            val func = firstArg<suspend (AmazonQLanguageServer) -> Unit>()
+            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() } just runs
+
+        sut.deleteTokenCredentials()
+
+        verify(exactly = 1) { mockLanguageServer.deleteTokenCredentials() }
+    }
+}