feat: support IAM role-based authentication for AWS Bedrock (#265)

Author: ssncferreira

config/config.go

@@ -26,6 +26,7 @@ type Anthropic struct {
 type AWSBedrock struct {
 	Region                     string
 	AccessKey, AccessKeySecret string
+	SessionToken               string
 	Model, SmallFastModel      string
 	// If set, requests will be sent to this URL instead of the default AWS Bedrock endpoint
 	// (https://bedrock-runtime.{region}.amazonaws.com).

intercept/messages/base.go

@@ -264,42 +264,60 @@ func (i *interceptionBase) withBody() option.RequestOption {
 	return option.WithRequestBody("application/json", []byte(i.reqPayload))
 }
 
+// withAWSBedrockOptions returns request options for authenticating with AWS Bedrock.
+//
+// When both AccessKey and AccessKeySecret are set in the aibridge config, they are
+// used directly as static credentials (with an optional SessionToken for temporary credentials).
+// Otherwise, the AWS SDK default credential chain resolves credentials (environment variables,
+// shared config/credentials files, IAM roles, IRSA, SSO, IMDS, etc.).
 func (*interceptionBase) withAWSBedrockOptions(ctx context.Context, cfg *aibconfig.AWSBedrock) ([]option.RequestOption, error) {
 	if cfg == nil {
 		return nil, xerrors.New("nil config given")
 	}
 	if cfg.Region == "" && cfg.BaseURL == "" {
 		return nil, xerrors.New("region or base url required")
 	}
-	if cfg.AccessKey == "" {
-		return nil, xerrors.New("access key required")
-	}
-	if cfg.AccessKeySecret == "" {
-		return nil, xerrors.New("access key secret required")
-	}
 	if cfg.Model == "" {
 		return nil, xerrors.New("model required")
 	}
 	if cfg.SmallFastModel == "" {
 		return nil, xerrors.New("small fast model required")
 	}
 
-	opts := []func(*config.LoadOptions) error{
+	loadOpts := []func(*config.LoadOptions) error{
 		config.WithRegion(cfg.Region),
-		config.WithCredentialsProvider(
+	}
+
+	// Use static credentials when explicitly provided, otherwise fall back to the SDK default credential chain.
+	switch {
+	// Both set: use static credentials directly.
+	case cfg.AccessKey != "" && cfg.AccessKeySecret != "":
+		loadOpts = append(loadOpts, config.WithCredentialsProvider(
 			credentials.NewStaticCredentialsProvider(
 				cfg.AccessKey,
 				cfg.AccessKeySecret,
-				"",
+				cfg.SessionToken, // optional
 			),
-		),
+		))
+	// Only one set: misconfiguration.
+	case cfg.AccessKey != "" || cfg.AccessKeySecret != "":
+		return nil, xerrors.New("both access key and access key secret must be provided together")
+	// Neither set: SDK default credential chain resolves credentials.
+	default:
 	}
 
-	awsCfg, err := config.LoadDefaultConfig(ctx, opts...)
+	awsCfg, err := config.LoadDefaultConfig(ctx, loadOpts...)
 	if err != nil {
 		return nil, xerrors.Errorf("failed to load AWS Bedrock config: %w", err)
 	}
 
+	// Fail fast: ensure credentials can be resolved before making any requests.
+	// awsCfg already carries the credentials provider, and the Bedrock middleware
+	// will call Retrieve on it when signing each request.
+	if _, err := awsCfg.Credentials.Retrieve(ctx); err != nil {
+		return nil, xerrors.Errorf("no AWS credentials found: %w", err)
+	}
+
 	var out []option.RequestOption
 	out = append(out, bedrock.WithConfig(awsCfg))
 

intercept/messages/base_test.go

@@ -84,9 +84,9 @@ func TestAWSBedrockValidation(t *testing.T) {
 		expectError bool
 		errorMsg    string
 	}{
-		// Valid cases.
+		// Valid cases: static credentials.
 		{
-			name: "valid with region",
+			name: "static credentials with region",
 			cfg: &config.AWSBedrock{
 				Region:          "us-east-1",
 				AccessKey:       "test-key",
@@ -96,7 +96,7 @@ func TestAWSBedrockValidation(t *testing.T) {
 			},
 		},
 		{
-			name: "valid with base url",
+			name: "static credentials with base url",
 			cfg: &config.AWSBedrock{
 				BaseURL:         "http://bedrock.internal",
 				AccessKey:       "test-key",
@@ -111,7 +111,7 @@ func TestAWSBedrockValidation(t *testing.T) {
 			// which is internal to the anthropic SDK.
 			//
 			// See TestAWSBedrockIntegration which validates this.
-			name: "valid with base url & region",
+			name: "static credentials with base url & region",
 			cfg: &config.AWSBedrock{
 				Region:          "us-east-1",
 				AccessKey:       "test-key",
@@ -120,6 +120,17 @@ func TestAWSBedrockValidation(t *testing.T) {
 				SmallFastModel:  "test-small-model",
 			},
 		},
+		{
+			name: "static credentials with session token",
+			cfg: &config.AWSBedrock{
+				Region:          "us-east-1",
+				AccessKey:       "test-key",
+				AccessKeySecret: "test-secret",
+				SessionToken:    "test-session-token",
+				Model:           "test-model",
+				SmallFastModel:  "test-small-model",
+			},
+		},
 		// Invalid cases.
 		{
 			name: "missing region & base url",
@@ -137,13 +148,12 @@ func TestAWSBedrockValidation(t *testing.T) {
 			name: "missing access key",
 			cfg: &config.AWSBedrock{
 				Region:          "us-east-1",
-				AccessKey:       "",
 				AccessKeySecret: "test-secret",
 				Model:           "test-model",
 				SmallFastModel:  "test-small-model",
 			},
 			expectError: true,
-			errorMsg:    "access key required",
+			errorMsg:    "both access key and access key secret must be provided together",
 		},
 		{
 			name: "missing access key secret",
@@ -155,7 +165,7 @@ func TestAWSBedrockValidation(t *testing.T) {
 				SmallFastModel:  "test-small-model",
 			},
 			expectError: true,
-			errorMsg:    "access key secret required",
+			errorMsg:    "both access key and access key secret must be provided together",
 		},
 		{
 			name: "missing model",
@@ -213,6 +223,82 @@ func TestAWSBedrockValidation(t *testing.T) {
 	}
 }
 
+// TestAWSBedrockCredentialChain tests credential resolution via the AWS SDK default credential chain.
+// NOTE: Cannot use t.Parallel() here because subtests use t.Setenv which requires sequential execution.
+func TestAWSBedrockCredentialChain(t *testing.T) {
+	tests := []struct {
+		name        string
+		cfg         *config.AWSBedrock
+		envVars     map[string]string
+		expectError bool
+		errorMsg    string
+	}{
+		{
+			name: "temporary credentials via env",
+			cfg: &config.AWSBedrock{
+				Region:         "us-east-1",
+				Model:          "test-model",
+				SmallFastModel: "test-small-model",
+			},
+			envVars: map[string]string{
+				"AWS_ACCESS_KEY_ID":     "test-key",
+				"AWS_SECRET_ACCESS_KEY": "test-secret",
+			},
+		},
+		{
+			name: "temporary credentials with session token via env",
+			cfg: &config.AWSBedrock{
+				Region:         "us-east-1",
+				Model:          "test-model",
+				SmallFastModel: "test-small-model",
+			},
+			envVars: map[string]string{
+				"AWS_ACCESS_KEY_ID":     "test-key",
+				"AWS_SECRET_ACCESS_KEY": "test-secret",
+				"AWS_SESSION_TOKEN":     "test-session-token",
+			},
+		},
+		{
+			// When static credentials are not provided and no environment credentials are set,
+			// the SDK default credential chain fails to resolve credentials.
+			name: "error when no credential source is configured",
+			cfg: &config.AWSBedrock{
+				Region:         "us-east-1",
+				Model:          "test-model",
+				SmallFastModel: "test-small-model",
+			},
+			envVars: map[string]string{
+				"AWS_ACCESS_KEY_ID":           "",
+				"AWS_SECRET_ACCESS_KEY":       "",
+				"AWS_SESSION_TOKEN":           "",
+				"AWS_PROFILE":                 "",
+				"AWS_SHARED_CREDENTIALS_FILE": "/dev/null",
+				"AWS_CONFIG_FILE":             "/dev/null",
+			},
+			expectError: true,
+			errorMsg:    "no AWS credentials found",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			for key, val := range tt.envVars {
+				t.Setenv(key, val)
+			}
+			base := &interceptionBase{}
+			opts, err := base.withAWSBedrockOptions(context.Background(), tt.cfg)
+
+			if tt.expectError {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), tt.errorMsg)
+			} else {
+				require.NotEmpty(t, opts)
+				require.NoError(t, err)
+			}
+		})
+	}
+}
+
 func TestAccumulateUsage(t *testing.T) {
 	t.Parallel()