Skip to content

Commit 352f368

Browse files
samgst-amazonsamgst-amazon
committed
implement aws/credentials/token messages
1 parent c222125 commit 352f368

File tree

4 files changed

+174
-0
lines changed

4 files changed

+174
-0
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()
1923
}

‎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<Void>
12+
}

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

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
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 CompletableFuture<ResponseMessage>().also { completableFuture ->
30+
AmazonQLspService.executeIfRunning(project) { server ->
31+
server.updateTokenCredentials(payload)
32+
.whenComplete { response, throwable ->
33+
if (throwable != null) {
34+
completableFuture.completeExceptionally(throwable)
35+
} else {
36+
completableFuture.complete(response)
37+
}
38+
}
39+
} ?: completableFuture.completeExceptionally(IllegalStateException("LSP Server not running"))
40+
}
41+
}
42+
43+
override fun deleteTokenCredentials(): CompletableFuture<Void> {
44+
return CompletableFuture<Void>().also { completableFuture ->
45+
AmazonQLspService.executeIfRunning(project) { server ->
46+
server.deleteTokenCredentials()
47+
completableFuture.complete(null)
48+
} ?: completableFuture.completeExceptionally(IllegalStateException("LSP Server not running"))
49+
}
50+
}
51+
52+
private fun createUpdateCredentialsPayload(token: String): UpdateCredentialsPayload {
53+
return UpdateCredentialsPayload(
54+
data = encryptionManager.encrypt(
55+
UpdateCredentialsPayloadData(
56+
BearerCredentials(token)
57+
)
58+
),
59+
encrypted = true
60+
)
61+
}
62+
}

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

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

0 commit comments

Comments
 (0)