Skip to content

Commit 1fb3d3b

Browse files
samgst-amazonleigaol
samgst-amazon
authored and
leigaol
committed
feat(amazonq): Implement aws/credentials/token messages (aws#5410)
* implement aws/credentials/token messages * detekt * feedback
1 parent c7661c1 commit 1fb3d3b

File tree

5 files changed

+164
-6
lines changed

5 files changed

+164
-6
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package software.aws.toolkits.jetbrains.services.amazonq.lsp
55

66
import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
7+
import org.eclipse.lsp4j.jsonrpc.services.JsonNotification
78
import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
89
import org.eclipse.lsp4j.services.LanguageServer
910
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
@@ -16,4 +17,7 @@ import java.util.concurrent.CompletableFuture
1617
interface AmazonQLanguageServer : LanguageServer {
1718
@JsonRequest("aws/credentials/token/update")
1819
fun updateTokenCredentials(payload: UpdateCredentialsPayload): CompletableFuture<ResponseMessage>
20+
21+
@JsonNotification("aws/credentials/token/delete")
22+
fun deleteTokenCredentials(): CompletableFuture<Unit>
1923
}

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -154,28 +154,26 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
154154
instance = start()
155155
}
156156

157-
suspend fun execute(runnable: suspend (AmazonQLanguageServer) -> Unit) {
157+
suspend fun<T> execute(runnable: suspend (AmazonQLanguageServer) -> T): T {
158158
val lsp = withTimeout(10.seconds) {
159159
val holder = mutex.withLock { instance }.await()
160160
holder.initializer.join()
161161

162162
holder.languageServer
163163
}
164-
165-
runnable(lsp)
164+
return runnable(lsp)
166165
}
167166

168-
fun executeSync(runnable: suspend (AmazonQLanguageServer) -> Unit) {
167+
fun<T> executeSync(runnable: suspend (AmazonQLanguageServer) -> T): T =
169168
runBlocking(cs.coroutineContext) {
170169
execute(runnable)
171170
}
172-
}
173171

174172
companion object {
175173
private val LOG = getLogger<AmazonQLspService>()
176174
fun getInstance(project: Project) = project.service<AmazonQLspService>()
177175

178-
fun executeIfRunning(project: Project, runnable: (AmazonQLanguageServer) -> Unit) =
176+
fun <T> executeIfRunning(project: Project, runnable: (AmazonQLanguageServer) -> T): T? =
179177
project.serviceIfCreated<AmazonQLspService>()?.executeSync(runnable)
180178
}
181179
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.amazonq.lsp.auth
5+
6+
import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
7+
import java.util.concurrent.CompletableFuture
8+
9+
interface AuthCredentialsService {
10+
fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture<ResponseMessage>
11+
fun deleteTokenCredentials(): CompletableFuture<Unit>
12+
}

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

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.amazonq.lsp.auth
5+
6+
import com.intellij.openapi.project.Project
7+
import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
8+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
9+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
10+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.BearerCredentials
11+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
12+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayloadData
13+
import java.util.concurrent.CompletableFuture
14+
15+
class DefaultAuthCredentialsService(
16+
private val project: Project,
17+
private val encryptionManager: JwtEncryptionManager,
18+
) : AuthCredentialsService {
19+
20+
override fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture<ResponseMessage> {
21+
val token = if (encrypted) {
22+
encryptionManager.decrypt(accessToken)
23+
} else {
24+
accessToken
25+
}
26+
27+
val payload = createUpdateCredentialsPayload(token)
28+
29+
return AmazonQLspService.executeIfRunning(project) { server ->
30+
server.updateTokenCredentials(payload)
31+
} ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
32+
}
33+
34+
override fun deleteTokenCredentials(): CompletableFuture<Unit> =
35+
CompletableFuture<Unit>().also { completableFuture ->
36+
AmazonQLspService.executeIfRunning(project) { server ->
37+
server.deleteTokenCredentials()
38+
completableFuture.complete(null)
39+
} ?: completableFuture.completeExceptionally(IllegalStateException("LSP Server not running"))
40+
}
41+
42+
private fun createUpdateCredentialsPayload(token: String): UpdateCredentialsPayload =
43+
UpdateCredentialsPayload(
44+
data = encryptionManager.encrypt(
45+
UpdateCredentialsPayloadData(
46+
BearerCredentials(token)
47+
)
48+
),
49+
encrypted = true
50+
)
51+
}

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

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.amazonq.lsp.auth
5+
6+
import com.intellij.openapi.components.serviceIfCreated
7+
import com.intellij.openapi.project.Project
8+
import io.mockk.every
9+
import io.mockk.mockk
10+
import io.mockk.verify
11+
import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
12+
import org.junit.Before
13+
import org.junit.Test
14+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLanguageServer
15+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
16+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
17+
import java.util.concurrent.CompletableFuture
18+
19+
class DefaultAuthCredentialsServiceTest {
20+
private lateinit var project: Project
21+
private lateinit var mockLanguageServer: AmazonQLanguageServer
22+
private lateinit var mockEncryptionManager: JwtEncryptionManager
23+
private lateinit var sut: DefaultAuthCredentialsService
24+
25+
@Before
26+
fun setUp() {
27+
project = mockk<Project>()
28+
mockLanguageServer = mockk<AmazonQLanguageServer>()
29+
mockEncryptionManager = mockk<JwtEncryptionManager>()
30+
every { mockEncryptionManager.encrypt(any()) } returns "mock-encrypted-data"
31+
32+
// Mock the service methods on Project
33+
val mockLspService = mockk<AmazonQLspService>()
34+
every { project.getService(AmazonQLspService::class.java) } returns mockLspService
35+
every { project.serviceIfCreated<AmazonQLspService>() } returns mockLspService
36+
37+
// Mock the LSP service's executeSync method as a suspend function
38+
every {
39+
mockLspService.executeSync<CompletableFuture<ResponseMessage>>(any())
40+
} coAnswers {
41+
val func = firstArg<suspend (AmazonQLanguageServer) -> CompletableFuture<ResponseMessage>>()
42+
func.invoke(mockLanguageServer)
43+
}
44+
45+
sut = DefaultAuthCredentialsService(project, this.mockEncryptionManager)
46+
}
47+
48+
@Test
49+
fun `test updateTokenCredentials unencrypted success`() {
50+
val token = "unencryptedToken"
51+
val isEncrypted = false
52+
53+
every {
54+
mockLanguageServer.updateTokenCredentials(any())
55+
} returns CompletableFuture.completedFuture(ResponseMessage())
56+
57+
sut.updateTokenCredentials(token, isEncrypted)
58+
59+
verify(exactly = 0) {
60+
mockEncryptionManager.decrypt(any())
61+
}
62+
verify(exactly = 1) {
63+
mockLanguageServer.updateTokenCredentials(any())
64+
}
65+
}
66+
67+
@Test
68+
fun `test updateTokenCredentials encrypted success`() {
69+
val encryptedToken = "encryptedToken"
70+
val decryptedToken = "decryptedToken"
71+
val isEncrypted = true
72+
73+
every { mockEncryptionManager.decrypt(encryptedToken) } returns decryptedToken
74+
every { mockEncryptionManager.encrypt(any()) } returns "mock-encrypted-data"
75+
every {
76+
mockLanguageServer.updateTokenCredentials(any())
77+
} returns CompletableFuture.completedFuture(ResponseMessage())
78+
79+
sut.updateTokenCredentials(encryptedToken, isEncrypted)
80+
81+
verify(exactly = 1) { mockEncryptionManager.decrypt(encryptedToken) }
82+
verify(exactly = 1) { mockLanguageServer.updateTokenCredentials(any()) }
83+
}
84+
85+
@Test
86+
fun `test deleteTokenCredentials success`() {
87+
every { mockLanguageServer.deleteTokenCredentials() } returns CompletableFuture.completedFuture(Unit)
88+
89+
sut.deleteTokenCredentials()
90+
91+
verify(exactly = 1) { mockLanguageServer.deleteTokenCredentials() }
92+
}
93+
}

0 commit comments

Comments
 (0)