add logs of requestId for ListAvailableCustomizations calls (#3982)

Author: Will-ShaoHua

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt

@@ -25,7 +25,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableC
 import software.amazon.awssdk.services.codewhispererruntime.model.OptOutPreference
 import software.amazon.awssdk.services.codewhispererruntime.model.SendTelemetryEventResponse
 import software.amazon.awssdk.services.codewhispererruntime.model.SuggestionState
-import software.amazon.awssdk.services.codewhispererruntime.paginators.ListAvailableCustomizationsIterable
+import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.core.AwsClientManager
@@ -36,6 +36,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.codewhisperer.customization.CodeWhispererCustomization
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestContext
@@ -79,7 +80,7 @@ interface CodeWhispererClientAdaptor : Disposable {
         isSigv4: Boolean = shouldUseSigv4Client(project)
     ): ListCodeScanFindingsResponse
 
-    fun listAvailableCustomizations(): ListAvailableCustomizationsIterable
+    fun listAvailableCustomizations(): List<CodeWhispererCustomization>
 
     fun sendUserTriggerDecisionTelemetry(
         requestContext: RequestContext,
@@ -190,8 +191,24 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
         }
 
     // DO NOT directly use this method to fetch customizations, use wrapper [CodeWhispererModelConfigurator.listCustomization()] instead
-    override fun listAvailableCustomizations(): ListAvailableCustomizationsIterable =
+    override fun listAvailableCustomizations(): List<CodeWhispererCustomization> =
         bearerClient().listAvailableCustomizationsPaginator(ListAvailableCustomizationsRequest.builder().build())
+            .stream()
+            .toList()
+            .flatMap { resp ->
+                LOG.debug {
+                    "listAvailableCustomizations: requestId: ${resp.responseMetadata().requestId()}, customizations: ${
+                        resp.customizations().map { it.name() }
+                    }"
+                }
+                resp.customizations().map {
+                    CodeWhispererCustomization(
+                        arn = it.arn(),
+                        name = it.name(),
+                        description = it.description()
+                    )
+                }
+            }
 
     override fun sendUserTriggerDecisionTelemetry(
         requestContext: RequestContext,

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt

@@ -119,11 +119,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
         calculateIfIamIdentityCenterConnection(project) {
             // 1. invoke API and get result
             val listAvailableCustomizationsResult = try {
-                val result = CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations().customizations().map { custom ->
-                    CodeWhispererCustomization(arn = custom.arn(), name = custom.name(), description = custom.description())
-                }
-                LOG.debug { "ListAvailableCustomizations: connection ${it.id} result: $result" }
-                result
+                CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations()
             } catch (e: Exception) {
                 val requestId = (e as? CodeWhispererRuntimeException)?.requestId()
                 val logMessage = if (CodeWhispererConstants.Customization.noAccessToCustomizationExceptionPredicate(e)) {

jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt

@@ -40,9 +40,12 @@ import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntime
 import software.amazon.awssdk.services.codewhispererruntime.model.CodeAnalysisStatus
 import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlRequest
 import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse
+import software.amazon.awssdk.services.codewhispererruntime.model.Customization
 import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
 import software.amazon.awssdk.services.codewhispererruntime.model.GetCodeAnalysisRequest
 import software.amazon.awssdk.services.codewhispererruntime.model.GetCodeAnalysisResponse
+import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableCustomizationsRequest
+import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableCustomizationsResponse
 import software.amazon.awssdk.services.codewhispererruntime.model.ListCodeAnalysisFindingsRequest
 import software.amazon.awssdk.services.codewhispererruntime.model.ListCodeAnalysisFindingsResponse
 import software.amazon.awssdk.services.codewhispererruntime.model.OptOutPreference
@@ -51,6 +54,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.SendTelemetryE
 import software.amazon.awssdk.services.codewhispererruntime.model.StartCodeAnalysisRequest
 import software.amazon.awssdk.services.codewhispererruntime.model.StartCodeAnalysisResponse
 import software.amazon.awssdk.services.codewhispererruntime.paginators.GenerateCompletionsIterable
+import software.amazon.awssdk.services.codewhispererruntime.paginators.ListAvailableCustomizationsIterable
 import software.amazon.awssdk.services.ssooidc.SsoOidcClient
 import software.aws.toolkits.core.TokenConnectionSettings
 import software.aws.toolkits.core.utils.test.aString
@@ -68,6 +72,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestU
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.sdkHttpResponse
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptorImpl
+import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
 import software.aws.toolkits.jetbrains.settings.AwsSettings
 
 class CodeWhispererClientAdaptorTest {
@@ -138,6 +143,48 @@ class CodeWhispererClientAdaptorTest {
         assertThat("us-east-1").isEqualTo(SONO_REGION)
     }
 
+    @Test
+    fun `listCustomizations`() {
+        val sdkIterable = ListAvailableCustomizationsIterable(bearerClient, ListAvailableCustomizationsRequest.builder().build())
+        val mockResponse1 = ListAvailableCustomizationsResponse.builder()
+            .customizations(
+                listOf(
+                    Customization.builder().name("custom-1").arn("arn-1").build(),
+                    Customization.builder().name("custom-2").arn("arn-2").build()
+                )
+            )
+            .nextToken("token-1")
+            .responseMetadata(metadata)
+            .sdkHttpResponse(sdkHttpResponse)
+            .build() as ListAvailableCustomizationsResponse
+
+        val mockResponse2 = ListAvailableCustomizationsResponse.builder()
+            .customizations(
+                listOf(
+                    Customization.builder().name("custom-3").arn("arn-3").build(),
+                )
+            )
+            .nextToken("")
+            .responseMetadata(metadata)
+            .sdkHttpResponse(sdkHttpResponse)
+            .build() as ListAvailableCustomizationsResponse
+
+        bearerClient.stub { client ->
+            on { client.listAvailableCustomizations(any<ListAvailableCustomizationsRequest>()) } doReturnConsecutively listOf(mockResponse1, mockResponse2)
+            on { client.listAvailableCustomizationsPaginator(any<ListAvailableCustomizationsRequest>()) } doReturn sdkIterable
+        }
+
+        val actual = sut.listAvailableCustomizations()
+        assertThat(actual).hasSize(3)
+        assertThat(actual).isEqualTo(
+            listOf(
+                CodeWhispererCustomization(name = "custom-1", arn = "arn-1"),
+                CodeWhispererCustomization(name = "custom-2", arn = "arn-2"),
+                CodeWhispererCustomization(name = "custom-3", arn = "arn-3")
+            )
+        )
+    }
+
     @Test
     fun `generateCompletionsPaginator - bearer`() {
         val request = pythonRequest