Skip to content

Commit 998db25

Browse files
Merge main into feature/dev-execution
2 parents 251d494 + fd7eb31 commit 998db25

File tree

3 files changed

+58
-57
lines changed

3 files changed

+58
-57
lines changed

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt

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

66
import com.intellij.openapi.Disposable
7+
import com.intellij.openapi.application.ApplicationManager
78
import com.intellij.openapi.components.service
89
import com.intellij.openapi.project.Project
910
import com.intellij.util.text.nullize
@@ -40,10 +41,14 @@ import software.amazon.awssdk.services.codewhispererruntime.model.TargetCode
4041
import software.amazon.awssdk.services.codewhispererruntime.model.UserIntent
4142
import software.aws.toolkits.core.utils.debug
4243
import software.aws.toolkits.core.utils.getLogger
44+
import software.aws.toolkits.core.utils.warn
4345
import software.aws.toolkits.jetbrains.core.AwsClientManager
4446
import software.aws.toolkits.jetbrains.core.awsClient
47+
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
48+
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
4549
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
46-
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
50+
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
51+
import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
4752
import software.aws.toolkits.jetbrains.services.amazonq.codeWhispererUserContext
4853
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
4954
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
@@ -62,6 +67,7 @@ import java.util.concurrent.TimeUnit
6267
import kotlin.reflect.KProperty0
6368
import kotlin.reflect.jvm.isAccessible
6469

70+
// TODO: move this file to package "/client"
6571
// As the connection is project-level, we need to make this project-level too
6672
@Deprecated("Methods can throw a NullPointerException if callee does not check if connection is valid")
6773
interface CodeWhispererClientAdaptor : Disposable {
@@ -277,16 +283,37 @@ interface CodeWhispererClientAdaptor : Disposable {
277283
open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispererClientAdaptor {
278284
private val mySigv4Client by lazy { createUnmanagedSigv4Client() }
279285

286+
@Volatile
287+
private var myBearerClient: CodeWhispererRuntimeClient? = null
288+
280289
private val KProperty0<*>.isLazyInitialized: Boolean
281290
get() {
282291
isAccessible = true
283292
return (getDelegate() as Lazy<*>).isInitialized()
284293
}
285294

286-
fun bearerClient(): CodeWhispererRuntimeClient =
287-
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.getConnectionSettings()
288-
?.awsClient<CodeWhispererRuntimeClient>()
289-
?: throw Exception("attempt to get bearer client while there is no valid credential")
295+
init {
296+
initClientUpdateListener()
297+
}
298+
299+
private fun initClientUpdateListener() {
300+
ApplicationManager.getApplication().messageBus.connect(this).subscribe(
301+
ToolkitConnectionManagerListener.TOPIC,
302+
object : ToolkitConnectionManagerListener {
303+
override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
304+
if (newConnection is AwsBearerTokenConnection) {
305+
myBearerClient = getBearerClient(newConnection.getConnectionSettings().providerId)
306+
}
307+
}
308+
}
309+
)
310+
}
311+
312+
private fun bearerClient(): CodeWhispererRuntimeClient {
313+
if (myBearerClient != null) return myBearerClient as CodeWhispererRuntimeClient
314+
myBearerClient = getBearerClient()
315+
return myBearerClient as CodeWhispererRuntimeClient
316+
}
290317

291318
override fun generateCompletionsPaginator(firstRequest: GenerateCompletionsRequest) = sequence<GenerateCompletionsResponse> {
292319
var nextToken: String? = firstRequest.nextToken()
@@ -827,6 +854,28 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
827854
if (this::mySigv4Client.isLazyInitialized) {
828855
mySigv4Client.close()
829856
}
857+
myBearerClient?.close()
858+
}
859+
860+
/**
861+
* Every different SSO/AWS Builder ID connection requires a new client which has its corresponding bearer token provider,
862+
* thus we have to create them dynamically.
863+
* Invalidate and recycle the old client first, and create a new client with the new connection.
864+
* This makes sure when we invoke CW, we always use the up-to-date connection.
865+
* In case this fails to close the client, myBearerClient is already set to null thus next time when we invoke CW,
866+
* it will go through this again which should get the current up-to-date connection. This stale client would be
867+
* unused and stay in memory for a while until eventually closed by ToolkitClientManager.
868+
*/
869+
open fun getBearerClient(oldProviderIdToRemove: String = ""): CodeWhispererRuntimeClient? {
870+
myBearerClient = null
871+
872+
val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())
873+
connection as? AwsBearerTokenConnection ?: run {
874+
LOG.warn { "$connection is not a bearer token connection" }
875+
return null
876+
}
877+
878+
return AwsClientManager.getInstance().getClient<CodeWhispererRuntimeClient>(connection.getConnectionSettings())
830879
}
831880

832881
companion object {
@@ -840,6 +889,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
840889
}
841890

842891
class MockCodeWhispererClientAdaptor(override val project: Project) : CodeWhispererClientAdaptorImpl(project) {
892+
override fun getBearerClient(oldProviderIdToRemove: String): CodeWhispererRuntimeClient = project.awsClient()
843893
override fun dispose() {}
844894
}
845895

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt

Lines changed: 2 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ import org.junit.After
1515
import org.junit.Before
1616
import org.junit.Rule
1717
import org.junit.Test
18-
import org.junit.jupiter.api.assertThrows
1918
import org.mockito.kotlin.any
2019
import org.mockito.kotlin.argThat
2120
import org.mockito.kotlin.argumentCaptor
@@ -69,20 +68,17 @@ import software.aws.toolkits.core.TokenConnectionSettings
6968
import software.aws.toolkits.core.utils.test.aString
7069
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
7170
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
72-
import software.aws.toolkits.jetbrains.core.credentials.DefaultToolkitConnectionManager
7371
import software.aws.toolkits.jetbrains.core.credentials.ManagedSsoProfile
7472
import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRule
7573
import software.aws.toolkits.jetbrains.core.credentials.MockToolkitAuthManagerRule
7674
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
77-
import software.aws.toolkits.jetbrains.core.credentials.logoutFromSsoConnection
78-
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
79-
import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES
8075
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_REGION
8176
import software.aws.toolkits.jetbrains.services.amazonq.FEATURE_EVALUATION_PRODUCT_NAME
8277
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata
8378
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonRequest
8479
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponseWithToken
8580
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.sdkHttpResponse
81+
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
8682
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptorImpl
8783
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
8884
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
@@ -113,7 +109,7 @@ class CodeWhispererClientAdaptorTest {
113109
private lateinit var bearerClient: CodeWhispererRuntimeClient
114110
private lateinit var ssoClient: SsoOidcClient
115111

116-
private lateinit var sut: CodeWhispererClientAdaptorImpl
112+
private lateinit var sut: CodeWhispererClientAdaptor
117113
private lateinit var connectionManager: ToolkitConnectionManager
118114
private var isTelemetryEnabledDefault: Boolean = false
119115

@@ -167,41 +163,6 @@ class CodeWhispererClientAdaptorTest {
167163
assertThat("us-east-1").isEqualTo(SONO_REGION)
168164
}
169165

170-
@Test
171-
fun `should throw if there is no valid credential, otherwise return codewhispererRuntimeClient`() {
172-
val connectionManager = DefaultToolkitConnectionManager()
173-
projectRule.project.replaceService(ToolkitConnectionManager::class.java, DefaultToolkitConnectionManager(), disposableRule.disposable)
174-
175-
assertThat(ToolkitConnectionManager.getInstance(projectRule.project).activeConnectionForFeature(QConnection.getInstance())).isNull()
176-
assertThrows<Exception>("attempt to get bearer client while there is no valid credential") {
177-
sut.bearerClient()
178-
}
179-
180-
val qConnection = authManagerRule.createConnection(ManagedSsoProfile("us-east-1", aString(), Q_SCOPES))
181-
connectionManager.switchConnection(qConnection)
182-
assertThat(connectionManager.activeConnectionForFeature(QConnection.getInstance()))
183-
.isNotNull
184-
.isEqualTo(qConnection)
185-
assertThat(sut.bearerClient())
186-
.isNotNull
187-
.isInstanceOf(CodeWhispererRuntimeClient::class.java)
188-
189-
logoutFromSsoConnection(projectRule.project, qConnection as AwsBearerTokenConnection)
190-
assertThat(connectionManager.activeConnectionForFeature(QConnection.getInstance())).isNull()
191-
assertThrows<Exception>("attempt to get bearer client while there is no valid credential") {
192-
sut.bearerClient()
193-
}
194-
195-
val anotherQConnection = authManagerRule.createConnection(ManagedSsoProfile("us-east-1", aString(), Q_SCOPES))
196-
connectionManager.switchConnection(anotherQConnection)
197-
assertThat(connectionManager.activeConnectionForFeature(QConnection.getInstance()))
198-
.isNotNull
199-
.isEqualTo(anotherQConnection)
200-
assertThat(sut.bearerClient())
201-
.isNotNull
202-
.isInstanceOf(CodeWhispererRuntimeClient::class.java)
203-
}
204-
205166
@Test
206167
fun `listCustomizations`() {
207168
val sdkIterable = ListAvailableCustomizationsIterable(bearerClient, ListAvailableCustomizationsRequest.builder().build())

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt

Lines changed: 1 addition & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -30,13 +30,8 @@ import org.mockito.kotlin.verify
3030
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
3131
import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
3232
import software.amazon.awssdk.services.codewhispererruntime.paginators.GenerateCompletionsIterable
33-
import software.amazon.awssdk.services.ssooidc.SsoOidcClient
3433
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
35-
import software.aws.toolkits.jetbrains.core.credentials.ManagedSsoProfile
3634
import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRule
37-
import software.aws.toolkits.jetbrains.core.credentials.MockToolkitAuthManagerRule
38-
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
39-
import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES
4035
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.codeWhispererRecommendationActionId
4136
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonFileName
4237
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponse
@@ -70,11 +65,10 @@ open class CodeWhispererTestBase {
7065
val mockClientManagerRule = MockClientManagerRule()
7166
val mockCredentialRule = MockCredentialManagerRule()
7267
val disposableRule = DisposableRule()
73-
val authManagerRule = MockToolkitAuthManagerRule()
7468

7569
@Rule
7670
@JvmField
77-
val ruleChain = RuleChain(projectRule, mockCredentialRule, mockClientManagerRule, authManagerRule, disposableRule)
71+
val ruleChain = RuleChain(projectRule, mockCredentialRule, mockClientManagerRule, disposableRule)
7872

7973
protected lateinit var mockClient: CodeWhispererRuntimeClient
8074

@@ -92,7 +86,6 @@ open class CodeWhispererTestBase {
9286
@Before
9387
open fun setUp() {
9488
mockClient = mockClientManagerRule.create()
95-
mockClientManagerRule.create<SsoOidcClient>()
9689
val requestCaptor = argumentCaptor<GenerateCompletionsRequest>()
9790
mockClient.stub {
9891
on {
@@ -166,9 +159,6 @@ open class CodeWhispererTestBase {
166159
projectRule.project.replaceService(CodeWhispererClientAdaptor::class.java, clientAdaptorSpy, disposableRule.disposable)
167160
ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, stateManager, disposableRule.disposable)
168161
stateManager.setAutoEnabled(false)
169-
170-
val conn = authManagerRule.createConnection(ManagedSsoProfile("us-east-1", "url", Q_SCOPES))
171-
ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(conn)
172162
}
173163

174164
@After

0 commit comments

Comments
 (0)