feat: emit bearer business metric (#1658)

* refactor tests and enable sourcing token from sysprops

* refactor

* consistent naming

* naming

* align with new constructor

* add kdoc

* rename transform file

* refactor transformers

* bump smithy version

Author: xinsong-cui

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/SdkIdTransform.kt

@@ -64,15 +64,15 @@ internal data class EndpointUrlConfigNames(
 
 internal fun String.toEndpointUrlConfigNames(): EndpointUrlConfigNames = EndpointUrlConfigNames(
     withTransform(JvmSystemPropertySuffix),
-    withTransform(SdkIdTransform.UpperSnakeCase),
-    withTransform(SdkIdTransform.LowerSnakeCase),
+    withTransform(SdkIdTransformers.UpperSnakeCase),
+    withTransform(SdkIdTransformers.LowerSnakeCase),
 )
 
 // JVM system property names follow the pattern "aws.endpointUrl${BaseClientName}"
 // where BaseClientName is the PascalCased sdk ID with any forbidden suffixes dropped - this is the same as what we use
 // for our client names
 // e.g. sdkId "Elasticsearch Service" -> client name "ElasticsearchClient", prop "aws.endpointUrlElasticsearch"
-private object JvmSystemPropertySuffix : SdkIdTransformer {
+private object JvmSystemPropertySuffix : StringTransformer {
     override fun transform(id: String): String =
         id.toPascalCase().removeSuffix("Service")
 }

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/ServiceClientCompanionObjectWriter.kt

@@ -0,0 +1,58 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.codegen
+
+import software.amazon.smithy.kotlin.codegen.utils.toPascalCase
+
+private val whitespaceRegex = Regex("\\s")
+
+private val spaceOrDashRegex = Regex("\\s|-")
+
+/**
+ * Base interface for string transformers.
+ */
+fun interface StringTransformer {
+    fun transform(input: String): String
+}
+
+/**
+ * Implements all standardized sdkId transforms.
+ */
+object SdkIdTransformers {
+    // Replace all whitespace from the sdkId with underscores and lowercase all letters.
+    val LowerSnakeCase = StringTransformer { it.replaceWhitespace("_").lowercase() }
+
+    // None. Directly use the sdkId.
+    val Identity = StringTransformer { it }
+
+    // Remove all whitespace from the sdkId.
+    val NoWhitespace = StringTransformer { it.replaceWhitespace("") }
+
+    // Replace all whitespace from the sdkId with dashes and lowercase all letters.
+    val LowerKebabCase = StringTransformer { it.replaceWhitespace("-").lowercase() }
+
+    // Replace all whitespace from the sdkId with underscores and capitalize all letters.
+    val UpperSnakeCase = StringTransformer { it.replaceWhitespace("_").uppercase() }
+}
+
+/**
+ * Implements all standardized SigV4 service signing name transforms.
+ */
+object SigV4NameTransformers {
+    // Replace all dashes from the SigV4 service signing name with underscores and capitalize all letters.
+    val UpperSnakeCase = StringTransformer { it.replaceSpaceOrDash("_").uppercase() }
+
+    // Remove dashes and convert SigV4 service signing name to PascalCase
+    val PascalCase = StringTransformer { it.toPascalCase() }
+}
+
+/**
+ * Applies the given transformer to the string.
+ */
+fun <T : StringTransformer> String.withTransform(transformer: T): String = transformer.transform(this)
+
+private fun String.replaceWhitespace(replacement: String) = replace(whitespaceRegex, replacement)
+
+private fun String.replaceSpaceOrDash(replacement: String) = replace(spaceOrDashRegex, replacement)

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/StringTransform.kt

@@ -4,8 +4,8 @@
  */
 package aws.sdk.kotlin.codegen.customization
 
-import aws.sdk.kotlin.codegen.SdkIdTransform
 import aws.sdk.kotlin.codegen.ServiceClientCompanionObjectWriter
+import aws.sdk.kotlin.codegen.SigV4NameTransformers
 import aws.sdk.kotlin.codegen.withTransform
 import software.amazon.smithy.kotlin.codegen.KotlinSettings
 import software.amazon.smithy.kotlin.codegen.core.*
@@ -22,15 +22,15 @@ import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.traits.HttpBearerAuthTrait
 
 /**
- * Customization that enables sourcing Bearer tokens from an environment variable
+ * Customization that enables sourcing Bearer tokens from JVM system properties and system environment variables
  *
- * When a service-specific environment variable for bearer tokens is present (e.g., AWS_BEARER_TOKEN_BEDROCK),
- * this customization configures the auth scheme resolver to prefer the smithy.api#httpBearerAuth scheme
- * over other authentication methods. Additionally, it configures a token provider that extracts the bearer token
- * from the target environment variable.
+ * When a service-specific JVM system property (e.g., aws.bearerTokenBedrock) or system environment variable
+ * for bearer tokens is present (e.g., AWS_BEARER_TOKEN_BEDROCK), this customization configures the
+ * auth scheme resolver to prefer the smithy.api#httpBearerAuth scheme over other authentication methods.
+ * Additionally, it configures a token provider that extracts the bearer token from these sources.
  */
 class EnvironmentBearerTokenCustomization : KotlinIntegration {
-    // Currently only services with sigv4 service name 'bedrock' need this customization
+    // Currently only services with sigV4 service name 'bedrock' need this customization
     private val supportedSigningServiceNames = setOf("bedrock")
 
     override fun enabledForService(model: Model, settings: KotlinSettings): Boolean {
@@ -68,9 +68,11 @@ class EnvironmentBearerTokenCustomization : KotlinIntegration {
     ) {
         val serviceSymbol = ctx.symbolProvider.toSymbol(serviceShape)
         val signingServiceName = AwsSignatureVersion4.signingServiceName(serviceShape)
-        // Transform signing name to environment variable name
-        val envVarSuffix = signingServiceName.withTransform(SdkIdTransform.UpperSnakeCase)
-        val envVarName = "AWS_BEARER_TOKEN_$envVarSuffix"
+        // Transform signing service name to environment variable key suffix
+        val envSuffix = signingServiceName.withTransform(SigV4NameTransformers.UpperSnakeCase)
+        val sysPropSuffix = signingServiceName.withTransform(SigV4NameTransformers.PascalCase)
+        val envKey = "AWS_BEARER_TOKEN_$envSuffix"
+        val sysPropKey = "aws.bearerToken$sysPropSuffix"
         val authSchemeId = RuntimeTypes.Auth.Identity.AuthSchemeId
 
         writer.withBlock(
@@ -79,9 +81,8 @@ class EnvironmentBearerTokenCustomization : KotlinIntegration {
             serviceSymbol,
             RuntimeTypes.Core.Utils.PlatformProvider,
         ) {
-            // The customization do nothing if environment variable is not set
-            write("if (provider.getenv(#S) == null) { return }", envVarName)
-
+            // The customization does nothing if environment variable and JVM system property are not set
+            write("if (provider.getProperty(#S) == null && provider.getenv(#S) == null) return", sysPropKey, envKey)
             // Configure auth scheme preference if customer hasn't specify one
             write("builder.config.authSchemePreference = builder.config.authSchemePreference ?: listOf(#T.HttpBearer)", authSchemeId)
 
@@ -93,9 +94,10 @@ class EnvironmentBearerTokenCustomization : KotlinIntegration {
             write("builder.config.authSchemePreference = listOf(#1T.HttpBearer) + filteredSchemes", authSchemeId)
 
             write(
-                "builder.config.bearerTokenProvider = builder.config.bearerTokenProvider ?: #T(#S, provider)",
+                "builder.config.bearerTokenProvider = builder.config.bearerTokenProvider ?: #T(#S, #S, provider)",
                 RuntimeTypes.Auth.HttpAuth.EnvironmentBearerTokenProvider,
-                envVarName,
+                sysPropKey,
+                envKey,
             )
         }
     }

codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/EnvironmentBearerTokenCustomization.kt

@@ -12,8 +12,8 @@ atomicfu-version = "0.29.0"
 binary-compatibility-validator-version = "0.18.0"
 
 # smithy-kotlin codegen and runtime are versioned separately
-smithy-kotlin-runtime-version = "1.5.1"
-smithy-kotlin-codegen-version = "0.35.1"
+smithy-kotlin-runtime-version = "1.5.4"
+smithy-kotlin-codegen-version = "0.35.4"
 
 # codegen
 smithy-version = "1.60.2"

gradle/libs.versions.toml

@@ -15,13 +15,14 @@ import aws.smithy.kotlin.runtime.http.engine.HttpClientEngineBase
 import aws.smithy.kotlin.runtime.http.engine.HttpClientEngineConfig
 import aws.smithy.kotlin.runtime.http.request.HttpRequest
 import aws.smithy.kotlin.runtime.http.response.HttpResponse
-import aws.smithy.kotlin.runtime.io.use
 import aws.smithy.kotlin.runtime.operation.ExecutionContext
 import aws.smithy.kotlin.runtime.time.Instant
 import aws.smithy.kotlin.runtime.util.TestPlatformProvider
 import kotlinx.coroutines.test.runTest
+import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
+import kotlin.test.assertTrue
 
 class BedrockEnvironmentBearerTokenTest {
     private fun mockHttpClient(handler: (HttpRequest) -> HttpResponse): HttpClientEngine {
@@ -39,69 +40,129 @@ class BedrockEnvironmentBearerTokenTest {
         env = mapOf("AWS_BEARER_TOKEN_BEDROCK" to "bedrock-token"),
     )
 
+    @Test
     fun testAuthSchemePreferenceConfigured() = runTest {
-        val builder = BedrockClient.Builder()
         val expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer)
+        val builder = BedrockClient.Builder()
+
         finalizeBearerTokenConfig(builder, mockPlatformProvider)
+
         assertEquals(expectedAuthSchemePreference, builder.config.authSchemePreference)
     }
 
+    @Test
     fun testBearerAuthSchemePromotedToFirst() = runTest {
+        val expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer, AuthSchemeId.AwsSigV4)
         val builder = BedrockClient.Builder()
+
         builder.config.authSchemePreference = listOf(AuthSchemeId.AwsSigV4)
         finalizeBearerTokenConfig(builder, mockPlatformProvider)
-        val expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer, AuthSchemeId.AwsSigV4)
         assertEquals(expectedAuthSchemePreference, builder.config.authSchemePreference)
 
         builder.config.authSchemePreference = listOf(AuthSchemeId.AwsSigV4, AuthSchemeId.HttpBearer)
+        finalizeBearerTokenConfig(builder, mockPlatformProvider)
         assertEquals(expectedAuthSchemePreference, builder.config.authSchemePreference)
     }
 
+    @Test
     fun testBearerTokenProviderConfigured() = runTest {
         val builder = BedrockClient.Builder()
         finalizeBearerTokenConfig(builder, mockPlatformProvider)
+
         assertNotNull(builder.config.bearerTokenProvider)
         val token = builder.config.bearerTokenProvider!!.resolve()
         assertNotNull(token)
         assertEquals("bedrock-token", token.token)
     }
 
+    @Test
+    fun testBearerTokenSourcingPrecedence() = runTest {
+        val builder = BedrockClient.Builder()
+
+        finalizeBearerTokenConfig(
+            builder,
+            TestPlatformProvider(
+                env = mapOf("AWS_BEARER_TOKEN_BEDROCK" to "env-bedrock-token"),
+                props = mapOf("aws.bearerTokenBedrock" to "sys-props-bedrock-token"),
+            ),
+        )
+
+        val token = builder.config.bearerTokenProvider!!.resolve()
+        assertEquals("sys-props-bedrock-token", token.token)
+    }
+
+    @Test
     fun testExplicitProviderTakesPrecedence() = runTest {
         val builder = BedrockClient.Builder()
+
         builder.config.bearerTokenProvider = object : BearerTokenProvider {
             override suspend fun resolve(attributes: Attributes): BearerToken = object : BearerToken {
                 override val token: String = "different-bedrock-token"
                 override val attributes: Attributes = emptyAttributes()
                 override val expiration: Instant? = null
             }
         }
+
         finalizeBearerTokenConfig(builder, mockPlatformProvider)
+
         assertNotNull(builder.config.bearerTokenProvider)
         val token = builder.config.bearerTokenProvider!!.resolve()
         assertNotNull(token)
         assertEquals("different-bedrock-token", token.token)
     }
 
+    @Test
     fun testBearerTokenProviderFunctionality() = runTest {
         var capturedAuthHeader: String? = null
 
-        BedrockClient.fromEnvironment {
-            region = "us-west-2"
-            httpClient = mockHttpClient { request ->
-                // Capture the Authorization header
+        val builder = BedrockClient.Builder().apply {
+            config.region = "us-west-2"
+            config.httpClient = mockHttpClient { request ->
                 capturedAuthHeader = request.headers["Authorization"]
                 HttpResponse(
                     status = HttpStatusCode.OK,
                     headers = Headers.Empty,
                     body = HttpBody.Empty,
                 )
             }
-        }.use { client ->
-            // Make an api call to capture Authorization header
-            client.listFoundationModels()
+        }
+
+        finalizeBearerTokenConfig(builder, mockPlatformProvider)
 
-            assertNotNull(capturedAuthHeader)
-            assertEquals("Bearer bedrock-token", capturedAuthHeader)
+        val testClient = builder.build()
+        // Make an api call to capture Authorization header
+        testClient.listFoundationModels()
+
+        assertNotNull(capturedAuthHeader)
+        assertEquals("Bearer bedrock-token", capturedAuthHeader)
+    }
+
+    @Test
+    fun testBusinessMetricEmitted() = runTest {
+        var capturedUserAgent: String? = null
+
+        val builder = BedrockClient.Builder().apply {
+            config.region = "us-west-2"
+            config.httpClient = mockHttpClient { request ->
+                capturedUserAgent = request.headers["User-Agent"]
+                HttpResponse(
+                    status = HttpStatusCode.OK,
+                    headers = Headers.Empty,
+                    body = HttpBody.Empty,
+                )
+            }
         }
+
+        finalizeBearerTokenConfig(builder, mockPlatformProvider)
+
+        val testClient = builder.build()
+        // Make an api call to capture User-Agent header
+        testClient.listFoundationModels()
+
+        assertNotNull(capturedUserAgent)
+        val capturedBusinessMetrics = Regex("m/([^\\s]+)").find(capturedUserAgent!!)?.value
+        assertNotNull(capturedBusinessMetrics)
+        // Check User-Agent header contains the business metric
+        assertTrue(capturedBusinessMetrics.contains("3"))
     }
 }