diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt index ffb1d909663..f18b7d0df75 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt @@ -7,7 +7,6 @@ import com.intellij.openapi.Disposable import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.components.service import com.intellij.openapi.project.Project -import com.intellij.openapi.util.SystemInfo import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider import software.amazon.awssdk.services.codewhisperer.CodeWhispererClient import software.amazon.awssdk.services.codewhisperer.model.CreateCodeScanRequest @@ -25,13 +24,10 @@ import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUr import software.amazon.awssdk.services.codewhispererruntime.model.Dimension import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse -import software.amazon.awssdk.services.codewhispererruntime.model.IdeCategory import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableCustomizationsRequest import software.amazon.awssdk.services.codewhispererruntime.model.ListFeatureEvaluationsResponse -import software.amazon.awssdk.services.codewhispererruntime.model.OperatingSystem import software.amazon.awssdk.services.codewhispererruntime.model.SendTelemetryEventResponse import software.amazon.awssdk.services.codewhispererruntime.model.SuggestionState -import software.amazon.awssdk.services.codewhispererruntime.model.UserContext import software.amazon.awssdk.services.codewhispererruntime.model.UserIntent import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger @@ -43,6 +39,7 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection +import software.aws.toolkits.jetbrains.services.amazonq.codeWhispererUserContext import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage @@ -51,10 +48,8 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestCon import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestContextNew import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants -import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.FEATURE_EVALUATION_PRODUCT_NAME import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getTelemetryOptOutPreference import software.aws.toolkits.jetbrains.services.codewhisperer.util.transform -import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata import software.aws.toolkits.telemetry.CodewhispererCompletionType import software.aws.toolkits.telemetry.CodewhispererSuggestionState import java.time.Instant @@ -201,24 +196,6 @@ interface CodeWhispererClientAdaptor : Disposable { } open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeWhispererClientAdaptor { - private val codeWhispererUserContext = ClientMetadata.getDefault().let { - val osForCodeWhisperer: OperatingSystem = - when { - SystemInfo.isWindows -> OperatingSystem.WINDOWS - SystemInfo.isMac -> OperatingSystem.MAC - // For now, categorize everything else as "Linux" (Linux/FreeBSD/Solaris/etc) - else -> OperatingSystem.LINUX - } - - UserContext.builder() - .ideCategory(IdeCategory.JETBRAINS) - .operatingSystem(osForCodeWhisperer) - .product(FEATURE_EVALUATION_PRODUCT_NAME) - .clientId(it.clientId) - .ideVersion(it.awsVersion) - .build() - } - private val mySigv4Client by lazy { createUnmanagedSigv4Client() } @Volatile @@ -348,7 +325,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW } } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(codeWhispererUserContext) + requestBuilder.userContext(codeWhispererUserContext()) } } @@ -391,7 +368,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW } } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(codeWhispererUserContext) + requestBuilder.userContext(codeWhispererUserContext()) } } @@ -413,7 +390,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW } } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(codeWhispererUserContext) + requestBuilder.userContext(codeWhispererUserContext()) } override fun sendUserModificationTelemetry( @@ -440,7 +417,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW } } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(codeWhispererUserContext) + requestBuilder.userContext(codeWhispererUserContext()) } override fun sendCodeScanTelemetry( @@ -459,7 +436,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW } } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(codeWhispererUserContext) + requestBuilder.userContext(codeWhispererUserContext()) } override fun sendCodeScanRemediationTelemetry( language: CodeWhispererProgrammingLanguage?, @@ -489,11 +466,11 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW } } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(codeWhispererUserContext) + requestBuilder.userContext(codeWhispererUserContext()) } override fun listFeatureEvaluations(): ListFeatureEvaluationsResponse = bearerClient().listFeatureEvaluations { - it.userContext(codeWhispererUserContext) + it.userContext(codeWhispererUserContext()) } override fun sendMetricDataTelemetry(eventName: String, metadata: Map): SendTelemetryEventResponse = @@ -507,7 +484,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW } } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(codeWhispererUserContext) + requestBuilder.userContext(codeWhispererUserContext()) } override fun sendChatAddMessageTelemetry( @@ -547,7 +524,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW } } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(codeWhispererUserContext) + requestBuilder.userContext(codeWhispererUserContext()) } override fun sendChatInteractWithMessageTelemetry( @@ -576,7 +553,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW telemetryEventBuilder.chatInteractWithMessageEvent(event) } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(codeWhispererUserContext) + requestBuilder.userContext(codeWhispererUserContext()) } override fun sendChatUserModificationTelemetry( @@ -602,7 +579,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW } } requestBuilder.optOutPreference(getTelemetryOptOutPreference()) - requestBuilder.userContext(codeWhispererUserContext) + requestBuilder.userContext(codeWhispererUserContext()) } override fun dispose() { diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt index 1d2f4b2788d..978bb92f08f 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt @@ -20,10 +20,10 @@ import com.intellij.util.xmlb.annotations.Property import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger +import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererFeatureConfigService import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants -import software.aws.toolkits.jetbrains.services.codewhisperer.util.calculateIfIamIdentityCenterConnection import software.aws.toolkits.jetbrains.utils.notifyInfo import software.aws.toolkits.jetbrains.utils.notifyWarn import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererFeatureConfigService.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererFeatureConfigService.kt index 5da5937b9ac..ba87955b9d0 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererFeatureConfigService.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererFeatureConfigService.kt @@ -7,12 +7,15 @@ import com.intellij.openapi.components.Service import com.intellij.openapi.components.service import com.intellij.openapi.project.Project import com.intellij.util.concurrency.annotations.RequiresBackgroundThread +import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient import software.amazon.awssdk.services.codewhispererruntime.model.FeatureValue +import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableCustomizationsRequest import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger -import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor -import software.aws.toolkits.jetbrains.services.codewhisperer.util.calculateIfBIDConnection -import software.aws.toolkits.jetbrains.services.codewhisperer.util.calculateIfIamIdentityCenterConnection +import software.aws.toolkits.jetbrains.core.awsClient +import software.aws.toolkits.jetbrains.services.amazonq.calculateIfBIDConnection +import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection +import software.aws.toolkits.jetbrains.services.amazonq.codeWhispererUserContext import software.aws.toolkits.jetbrains.utils.isQExpired @Service @@ -25,7 +28,9 @@ class CodeWhispererFeatureConfigService { LOG.debug { "Fetching feature configs" } try { - val response = CodeWhispererClientAdaptor.getInstance(project).listFeatureEvaluations() + val response = project.awsClient().listFeatureEvaluations { + it.userContext(codeWhispererUserContext()) + } // Simply force overwrite feature configs from server response, no needed to check existing values. response.featureEvaluations().forEach { @@ -49,7 +54,16 @@ class CodeWhispererFeatureConfigService { val availableCustomizations = calculateIfIamIdentityCenterConnection(project) { try { - CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations().map { c -> c.arn } + project.awsClient().listAvailableCustomizationsPaginator( + ListAvailableCustomizationsRequest.builder().build() + ) + .stream() + .toList() + .flatMap { resp -> + resp.customizations().map { + it.arn() + } + } } catch (e: Exception) { LOG.debug(e) { "Failed to list available customizations" } null @@ -71,9 +85,11 @@ class CodeWhispererFeatureConfigService { } fun getFeatureConfigsTelemetry(): String = - "{${featureConfigs.entries.joinToString(", ") { (name, context) -> - "$name: ${context.variation}" - }}}" + "{${ + featureConfigs.entries.joinToString(", ") { (name, context) -> + "$name: ${context.variation}" + } + }}" // TODO: for all feature variations, define a contract that can be enforced upon the implementation of // the business logic. @@ -88,7 +104,8 @@ class CodeWhispererFeatureConfigService { // 6) Add a test case for this feature. fun getTestFeature(): String = getFeatureValueForKey(TEST_FEATURE_NAME).stringValue() - fun getIsDataCollectionEnabled(): Boolean = getFeatureValueForKey(DATA_COLLECTION_FEATURE).stringValue() == "data-collection" + fun getIsDataCollectionEnabled(): Boolean = + getFeatureValueForKey(DATA_COLLECTION_FEATURE).stringValue() == "data-collection" fun getCustomizationArnOverride(): String = getFeatureValueForKey(CUSTOMIZATION_ARN_OVERRIDE_NAME).stringValue() diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt index ab20e89b0b7..7cd8b6a83cf 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt @@ -10,6 +10,7 @@ import kotlinx.coroutines.delay import kotlinx.coroutines.isActive import kotlinx.coroutines.launch import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope +import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager @@ -20,7 +21,6 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispe import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.FEATURE_CONFIG_POLL_INTERVAL_IN_MS import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.promptReAuth -import software.aws.toolkits.jetbrains.services.codewhisperer.util.calculateIfIamIdentityCenterConnection import software.aws.toolkits.jetbrains.utils.isQConnected import software.aws.toolkits.jetbrains.utils.isQExpired import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt index ba1a3ba6851..08a9cdee985 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt @@ -31,7 +31,6 @@ object CodeWhispererConstants { const val POPUP_DELAY_CHECK_INTERVAL: Long = 25 const val IDLE_TIME_CHECK_INTERVAL: Long = 25 const val SUPPLEMENTAL_CONTEXT_TIMEOUT = 50L - const val FEATURE_EVALUATION_PRODUCT_NAME = "CodeWhisperer" val AWSTemplateKeyWordsRegex = Regex("(AWSTemplateFormatVersion|Resources|AWS::|Description)") val AWSTemplateCaseInsensitiveKeyWordsRegex = Regex("(cloudformation|cfn|template|description)") diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt index 78744266a42..deffcb7e6f5 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt @@ -50,27 +50,6 @@ import software.aws.toolkits.resources.message import software.aws.toolkits.telemetry.CodewhispererCompletionType import software.aws.toolkits.telemetry.CodewhispererGettingStartedTask -fun calculateIfIamIdentityCenterConnection(project: Project, calculationTask: (connection: ToolkitConnection) -> T): T? = - ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())?.let { - calculateIfIamIdentityCenterConnection(it, calculationTask) - } - -fun calculateIfIamIdentityCenterConnection(connection: ToolkitConnection, calculationTask: (connection: ToolkitConnection) -> T): T? = - if (connection.isSono()) { - null - } else { - calculationTask(connection) - } - -fun calculateIfBIDConnection(project: Project, calculationTask: (connection: ToolkitConnection) -> T): T? = - ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())?.let { - if (it.isSono()) { - calculationTask(it) - } else { - null - } - } - // Controls the condition to send telemetry event to CodeWhisperer service, currently: // 1. It will be sent for Builder ID users, only if they have optin telemetry sharing. // 2. It will be sent for IdC users, regardless of telemetry optout status. diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt index c98639cd067..ab884eef500 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt @@ -73,6 +73,7 @@ import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRul import software.aws.toolkits.jetbrains.core.credentials.MockToolkitAuthManagerRule import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_REGION +import software.aws.toolkits.jetbrains.services.amazonq.FEATURE_EVALUATION_PRODUCT_NAME import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonRequest import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponseWithToken @@ -87,7 +88,6 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispe import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants -import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.FEATURE_EVALUATION_PRODUCT_NAME import software.aws.toolkits.jetbrains.settings.AwsSettings import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule import software.aws.toolkits.telemetry.CodewhispererCompletionType diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt index 30ce6fc30df..cc60a04e19a 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt @@ -11,19 +11,25 @@ import kotlinx.coroutines.runBlocking import org.assertj.core.api.Assertions.assertThat import org.junit.Ignore import org.junit.Rule +import org.mockito.kotlin.any import org.mockito.kotlin.doReturn import org.mockito.kotlin.eq import org.mockito.kotlin.mock import org.mockito.kotlin.stub +import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient +import software.amazon.awssdk.services.codewhispererruntime.model.Customization import software.amazon.awssdk.services.codewhispererruntime.model.FeatureEvaluation import software.amazon.awssdk.services.codewhispererruntime.model.FeatureValue +import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableCustomizationsRequest +import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableCustomizationsResponse +import software.amazon.awssdk.services.codewhispererruntime.model.ListFeatureEvaluationsRequest import software.amazon.awssdk.services.codewhispererruntime.model.ListFeatureEvaluationsResponse +import software.amazon.awssdk.services.codewhispererruntime.paginators.ListAvailableCustomizationsIterable +import software.aws.toolkits.jetbrains.core.MockClientManagerRule import software.aws.toolkits.jetbrains.core.credentials.LegacyManagedBearerSsoConnection import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_URL -import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor -import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererFeatureConfigService import kotlin.reflect.full.memberFunctions import kotlin.test.Test @@ -41,6 +47,10 @@ class CodeWhispererFeatureConfigServiceTest { @Rule val projectRule = ProjectRule() + @JvmField + @Rule + val mockClientManagerRule = MockClientManagerRule() + @Test fun `test FEATURE_DEFINITIONS is not empty`() { assertThat(CodeWhispererFeatureConfigService.FEATURE_DEFINITIONS).isNotEmpty @@ -64,9 +74,8 @@ class CodeWhispererFeatureConfigServiceTest { } private fun testCustomizationArnOverrideABHelper(isIdc: Boolean, isInListAvailableCustomizations: Boolean) { - val clientAdaptorSpy = mock() - clientAdaptorSpy.stub { - on { listFeatureEvaluations() } doReturn ListFeatureEvaluationsResponse.builder().featureEvaluations( + mockClientManagerRule.create().stub { + on { listFeatureEvaluations(any()) } doReturn ListFeatureEvaluationsResponse.builder().featureEvaluations( listOf( FeatureEvaluation.builder() .feature(CodeWhispererFeatureConfigService.CUSTOMIZATION_ARN_OVERRIDE_NAME) @@ -75,12 +84,26 @@ class CodeWhispererFeatureConfigServiceTest { .build() ) ).build() - on { listAvailableCustomizations() } doReturn + + val mockResponseIterable: ListAvailableCustomizationsIterable = mock() + mockResponseIterable.stub { if (isInListAvailableCustomizations) { - listOf(CodeWhispererCustomization(arn = "test arn", name = "Test Arn")) + on { stream() } doReturn listOf( + ListAvailableCustomizationsResponse.builder() + .customizations( + Customization.builder().arn("test arn").name("Test Arn").build() + ).build() + ).stream() } else { - emptyList() + on { stream() } doReturn listOf( + ListAvailableCustomizationsResponse.builder() + .customizations( + emptyList() + ).build() + ).stream() } + } + on { listAvailableCustomizationsPaginator(any()) } doReturn mockResponseIterable } val mockSsoConnection = mock { @@ -93,12 +116,6 @@ class CodeWhispererFeatureConfigServiceTest { disposableRule.disposable ) - projectRule.project.replaceService( - CodeWhispererClientAdaptor::class.java, - clientAdaptorSpy, - disposableRule.disposable - ) - runBlocking { CodeWhispererFeatureConfigService.getInstance().fetchFeatureConfigs(projectRule.project) } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/Constants.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/Constants.kt index 4b298106c7b..aba815bdc78 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/Constants.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/Constants.kt @@ -35,3 +35,5 @@ const val CODE_TRANSFORM_TROUBLESHOOT_DOC_DOWNLOAD_ERROR_OVERVIEW = const val CODE_TRANSFORM_PREREQUISITES = "https://docs.aws.amazon.com/amazonq/latest/qdeveloper-ug/code-transformation.html#prerequisites" + +const val FEATURE_EVALUATION_PRODUCT_NAME = "CodeWhisperer" diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QUtils.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QUtils.kt new file mode 100644 index 00000000000..f6dea5c681e --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QUtils.kt @@ -0,0 +1,54 @@ +// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package software.aws.toolkits.jetbrains.services.amazonq + +import com.intellij.openapi.project.Project +import com.intellij.openapi.util.SystemInfo +import software.amazon.awssdk.services.codewhispererruntime.model.IdeCategory +import software.amazon.awssdk.services.codewhispererruntime.model.OperatingSystem +import software.amazon.awssdk.services.codewhispererruntime.model.UserContext +import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection +import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager +import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection +import software.aws.toolkits.jetbrains.core.credentials.sono.isSono +import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata + +fun calculateIfIamIdentityCenterConnection(project: Project, calculationTask: (connection: ToolkitConnection) -> T): T? = + ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())?.let { + calculateIfIamIdentityCenterConnection(it, calculationTask) + } + +fun calculateIfIamIdentityCenterConnection(connection: ToolkitConnection, calculationTask: (connection: ToolkitConnection) -> T): T? = + if (connection.isSono()) { + null + } else { + calculationTask(connection) + } + +fun calculateIfBIDConnection(project: Project, calculationTask: (connection: ToolkitConnection) -> T): T? = + ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance())?.let { + if (it.isSono()) { + calculationTask(it) + } else { + null + } + } + +fun codeWhispererUserContext(): UserContext = ClientMetadata.getDefault().let { + val osForCodeWhisperer: OperatingSystem = + when { + SystemInfo.isWindows -> OperatingSystem.WINDOWS + SystemInfo.isMac -> OperatingSystem.MAC + // For now, categorize everything else as "Linux" (Linux/FreeBSD/Solaris/etc) + else -> OperatingSystem.LINUX + } + + UserContext.builder() + .ideCategory(IdeCategory.JETBRAINS) + .operatingSystem(osForCodeWhisperer) + .product(FEATURE_EVALUATION_PRODUCT_NAME) + .clientId(it.clientId) + .ideVersion(it.awsVersion) + .build() +} diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomization.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomization.kt similarity index 100% rename from plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomization.kt rename to plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomization.kt