Support AWS_SESSION_TOKEN in BedrockClient

Author: tiulpin

.github/workflows/heavy-tests.yml

@@ -10,10 +10,11 @@ name: Heavy Tests
 on:
   workflow_dispatch:  # Manual trigger
   push:
-    branches: [ "main", "develop" ]
+    branches: [ main, develop ]
 
 env:
   AWS_REGION: us-west-2
+  KOOG_HEAVY_TESTS: true
 
 jobs:
   integration-tests:
@@ -52,6 +53,7 @@ jobs:
           role-to-assume: ${{ secrets.AWS_ROLE_TO_ASSUME_KOOG }}
           aws-region: ${{ env.AWS_REGION }}
           role-session-name: GitHub-Actions-Koog-Integration-Tests
+          role-duration-seconds: 10800  # 3 hours (3 * 60 * 60)
 
       - name: Verify AWS Identity and Bedrock access
         run: |
@@ -66,6 +68,12 @@ jobs:
           }
           echo "Bedrock access verified!"
 
+      - name: Run Bedrock Credentials Smoke Test
+        run: |
+          echo "Running Bedrock credentials smoke test..."
+          ./gradlew :integration-tests:cleanJvmTest :integration-tests:jvmTest --tests "ai.koog.integration.tests.BedrockCredentialsSmokeTest"
+          echo "Bedrock credentials smoke test passed!"
+
       - name: JvmIntegrationTest with Gradle Wrapper
         env:
           ANTHROPIC_API_TEST_KEY: ${{ secrets.ANTHROPIC_API_TEST_KEY }}

build.gradle.kts

@@ -90,7 +90,7 @@ subprojects {
                 "OPEN_AI_API_TEST_KEY" to System.getenv("OPEN_AI_API_TEST_KEY"),
                 "GEMINI_API_TEST_KEY" to System.getenv("GEMINI_API_TEST_KEY"),
                 "OPEN_ROUTER_API_TEST_KEY" to System.getenv("OPEN_ROUTER_API_TEST_KEY"),
-                "AWS_SECRET_KEY" to System.getenv("AWS_SECRET_KEY"),
+                "AWS_SECRET_ACCESS_KEY" to System.getenv("AWS_SECRET_ACCESS_KEY"),
                 "AWS_ACCESS_KEY_ID" to System.getenv("AWS_ACCESS_KEY_ID"),
             )
         )

gradle/libs.versions.toml

@@ -56,7 +56,9 @@ opentelemetry-bom = { module = "io.opentelemetry:opentelemetry-bom", version.ref
 opentelemetry-sdk = { module = "io.opentelemetry:opentelemetry-sdk" }
 opentelemetry-exporter-otlp = { module = "io.opentelemetry:opentelemetry-exporter-otlp" }
 opentelemetry-exporter-logging = { module = "io.opentelemetry:opentelemetry-exporter-logging" }
+aws-sdk-kotlin-bedrock = { module = "aws.sdk.kotlin:bedrock", version.ref = "aws-sdk-kotlin" }
 aws-sdk-kotlin-bedrockruntime = { module = "aws.sdk.kotlin:bedrockruntime", version.ref = "aws-sdk-kotlin" }
+aws-sdk-kotlin-sts = { module = "aws.sdk.kotlin:sts", version.ref = "aws-sdk-kotlin" }
 
 # Spring
 spring-boot-bom = { module = "org.springframework.boot:spring-boot-dependencies", version.ref = "spring-boot" }

integration-tests/build.gradle.kts

@@ -31,6 +31,9 @@ kotlin {
                 implementation(libs.junit.jupiter.params)
                 implementation(libs.kotlinx.coroutines.test)
                 implementation(libs.kotlinx.serialization.json)
+                implementation(libs.aws.sdk.kotlin.sts)
+                implementation(libs.aws.sdk.kotlin.bedrock)
+                implementation(libs.aws.sdk.kotlin.bedrockruntime)
                 implementation(libs.ktor.client.content.negotiation)
                 runtimeOnly(libs.slf4j.simple)
             }

integration-tests/env.template.properties

@@ -7,4 +7,9 @@
 OPEN_AI_API_TEST_KEY=
 ANTHROPIC_API_TEST_KEY=
 GEMINI_API_TEST_KEY=
-OPEN_ROUTER_API_TEST_KEY=
+OPEN_ROUTER_API_TEST_KEY=
+
+# AWS Bedrock credentials
+AWS_ACCESS_KEY_ID=
+AWS_SECRET_ACCESS_KEY=
+AWS_SESSION_TOKEN=

integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/BedrockCredentialsSmokeTest.kt

@@ -0,0 +1,34 @@
+package ai.koog.integration.tests
+
+import aws.sdk.kotlin.services.bedrock.BedrockClient
+import aws.sdk.kotlin.services.bedrock.listFoundationModels
+import kotlinx.coroutines.runBlocking
+import org.junit.jupiter.api.condition.DisabledOnOs
+import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable
+import org.junit.jupiter.api.condition.OS
+import kotlin.test.Test
+import kotlin.test.assertTrue
+
+/**
+ * Verifies that the credentials provided to the build can access the Bedrock control-plane
+ * API (ListFoundationModels). If the call fails we abort early instead of waiting for a
+ * full prompt execution to blow up.
+ *
+ * This test only runs when AWS credentials are available (typically in heavy-tests workflow).
+ */
+class BedrockCredentialsSmokeTest {
+    @Test
+    @EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".+")
+    @EnabledIfEnvironmentVariable(named = "KOOG_HEAVY_TESTS", matches = "true")
+    fun listFoundationModelsWorks() = runBlocking {
+        val region = System.getenv("AWS_REGION") ?: "us-west-2"
+
+        BedrockClient { this.region = region }.use { bedrock ->
+            val resp = bedrock.listFoundationModels { }
+            assertTrue(
+                (resp.modelSummaries?.isNotEmpty()) == true,
+                "Bedrock ListFoundationModels returned no models – credentials/region might be wrong"
+            )
+        }
+    }
+}

integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/SingleLLMPromptExecutorIntegrationTest.kt

@@ -16,6 +16,7 @@ import ai.koog.integration.tests.utils.RetryUtils.withRetry
 import ai.koog.integration.tests.utils.TestUtils
 import ai.koog.integration.tests.utils.TestUtils.readAwsAccessKeyIdFromEnv
 import ai.koog.integration.tests.utils.TestUtils.readAwsSecretAccessKeyFromEnv
+import ai.koog.integration.tests.utils.TestUtils.readAwsSessionTokenFromEnv
 import ai.koog.integration.tests.utils.TestUtils.readTestAnthropicKeyFromEnv
 import ai.koog.integration.tests.utils.TestUtils.readTestGoogleAIKeyFromEnv
 import ai.koog.integration.tests.utils.TestUtils.readTestOpenAIKeyFromEnv
@@ -44,6 +45,7 @@ import kotlinx.coroutines.test.runTest
 import org.junit.jupiter.api.Assertions.assertNotNull
 import org.junit.jupiter.api.Assumptions.assumeTrue
 import org.junit.jupiter.api.BeforeAll
+import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.Arguments
 import org.junit.jupiter.params.provider.MethodSource
@@ -79,6 +81,7 @@ class SingleLLMPromptExecutorIntegrationTest {
             val bedrockClientInstance = BedrockLLMClient(
                 readAwsAccessKeyIdFromEnv(),
                 readAwsSecretAccessKeyFromEnv(),
+                readAwsSessionTokenFromEnv(),
                 BedrockClientSettings()
             )
             // val openRouterClientInstance = OpenRouterLLMClient(readTestOpenRouterKeyFromEnv())
@@ -101,6 +104,7 @@ class SingleLLMPromptExecutorIntegrationTest {
             val bedrockClientInstance = BedrockLLMClient(
                 readAwsAccessKeyIdFromEnv(),
                 readAwsSecretAccessKeyFromEnv(),
+                readAwsSessionTokenFromEnv(),
                 BedrockClientSettings(),
             )
 
@@ -1044,7 +1048,8 @@ class SingleLLMPromptExecutorIntegrationTest {
     fun integration_testSimpleBedrockExecutor(model: LLModel) = runTest(timeout = 300.seconds) {
         val executor = simpleBedrockExecutor(
             readAwsAccessKeyIdFromEnv(),
-            readAwsSecretAccessKeyFromEnv()
+            readAwsSecretAccessKeyFromEnv(),
+            readAwsSessionTokenFromEnv() ?: "",
         )
 
         val prompt = Prompt.build("test-simple-bedrock-executor") {

integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/utils/Models.kt

@@ -76,7 +76,7 @@ object Models {
     @JvmStatic
     fun bedrockModels(): Stream<LLModel> {
         return Stream.of(
-            BedrockModels.AI21JambaMini,
+            // BedrockModels.AI21JambaMini, // Disabled: Only available in us-east-2, workflow uses us-west-2
             BedrockModels.AmazonNovaLite,
             BedrockModels.AnthropicClaude35Haiku,
             BedrockModels.MetaLlama3_1_8BInstruct,

integration-tests/src/jvmTest/kotlin/ai/koog/integration/tests/utils/TestUtils.kt

@@ -32,8 +32,15 @@ object TestUtils {
     }
 
     fun readAwsSecretAccessKeyFromEnv(): String {
-        return System.getenv("AWS_SECRET_KEY")
-            ?: error("ERROR: environment variable `AWS_SECRET_KEY` is not set")
+        return System.getenv("AWS_SECRET_ACCESS_KEY")
+            ?: error("ERROR: environment variable `AWS_SECRET_ACCESS_KEY` is not set")
+    }
+
+    fun readAwsSessionTokenFromEnv(): String? {
+        return System.getenv("AWS_SESSION_TOKEN")
+            ?: null.also {
+                println("WARNING: environment variable `AWS_SESSION_TOKEN` is not set, using default session token")
+            }
     }
 
     @Serializable

prompt/prompt-executor/prompt-executor-clients/prompt-executor-bedrock-client/src/jvmMain/kotlin/ai/koog/prompt/executor/clients/bedrock/BedrockLLMClient.kt

@@ -93,13 +93,15 @@ public class BedrockLLMClient(
      *
      * @param awsAccessKeyId The AWS access key ID for authentication
      * @param awsSecretAccessKey The AWS secret access key for authentication
+     * @param awsSessionToken Optional session token for temporary credentials
      * @param settings Configuration settings for the Bedrock client, such as region and endpoint
      * @param clock A clock used for time-based operations
      * @return A configured [LLMClient] instance for Bedrock
      */
     public constructor(
         awsAccessKeyId: String,
         awsSecretAccessKey: String,
+        awsSessionToken: String? = null,
         settings: BedrockClientSettings = BedrockClientSettings(),
         clock: Clock = Clock.System,
     ) : this(
@@ -108,6 +110,7 @@ public class BedrockLLMClient(
             this.credentialsProvider = StaticCredentialsProvider {
                 this.accessKeyId = awsAccessKeyId
                 this.secretAccessKey = awsSecretAccessKey
+                awsSessionToken?.let { this.sessionToken = it }
             }
 
             // Configure custom endpoint if provided