From bd4b82679cb53dbdd9e226e37ced8e8d323df471 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 2 Jul 2026 16:22:13 +0000 Subject: [PATCH 1/2] fix(aibridge/provider): disable keep-alive on the STS assume-role client --- aibridge/provider/bedrock.go | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/aibridge/provider/bedrock.go b/aibridge/provider/bedrock.go index 013430f1d7a..23f5b1db6aa 100644 --- a/aibridge/provider/bedrock.go +++ b/aibridge/provider/bedrock.go @@ -2,8 +2,10 @@ package provider import ( "context" + "net/http" "github.com/aws/aws-sdk-go-v2/aws" + awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http" awsconfig "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/credentials" "github.com/aws/aws-sdk-go-v2/credentials/stscreds" @@ -77,7 +79,24 @@ func buildBedrockCredentials(ctx context.Context, cfg config.AWSBedrock) (aws.Cr // cache to avoid re-assuming the role on every request. credsProvider := base.Credentials if cfg.RoleARN != "" { - credsProvider = stscreds.NewAssumeRoleProvider(sts.NewFromConfig(base), cfg.RoleARN, func(o *stscreds.AssumeRoleOptions) { + // Disable keep-alive on the STS client so each AssumeRole opens a + // fresh connection. Observed: with keep-alive, AssumeRole calls reuse + // one connection pinned to a single STS endpoint, and after a + // trust-policy change that connection kept returning AccessDenied for + // minutes while a fresh connection (e.g. the AWS CLI) accepted the + // identical request at once; the gateway recovered only when that + // connection recycled or the process restarted. The STS-internal reason is + // unconfirmed (likely per-endpoint propagation of the change); what we + // verified is that a fresh connection per call recovers in seconds + // instead of minutes. AssumeRole runs at most once per credential-cache + // lifetime, so keep-alive saves nothing here. Scoped to the STS client + // only; Bedrock requests use a separate client and keep pooling. + stsClient := sts.NewFromConfig(base, func(o *sts.Options) { + o.HTTPClient = awshttp.NewBuildableClient().WithTransportOptions(func(t *http.Transport) { + t.DisableKeepAlives = true + }) + }) + credsProvider = stscreds.NewAssumeRoleProvider(stsClient, cfg.RoleARN, func(o *stscreds.AssumeRoleOptions) { o.RoleSessionName = bedrockSessionName if cfg.ExternalID != "" { o.ExternalID = aws.String(cfg.ExternalID) From fa282e457b58fdf8cac98f3104ae20a7c4810eed Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 2 Jul 2026 18:44:46 +0000 Subject: [PATCH 2/2] test: assert the STS client disables keep-alive --- aibridge/provider/bedrock_internal_test.go | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/aibridge/provider/bedrock_internal_test.go b/aibridge/provider/bedrock_internal_test.go index 1cb7764c894..e9827b7895f 100644 --- a/aibridge/provider/bedrock_internal_test.go +++ b/aibridge/provider/bedrock_internal_test.go @@ -159,13 +159,15 @@ func TestBuildBedrockCredentialsDefaultChain(t *testing.T) { // name are sent and that the returned temporary credentials are used. // NOTE: no t.Parallel() because it uses t.Setenv. func TestBuildBedrockCredentialsAssumeRole(t *testing.T) { - var gotRoleARN, gotSessionName string + var gotRoleARN, gotSessionName, gotConnection string // Mock the AWS STS AssumeRole API. // https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html sts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.NoError(t, r.ParseForm()) gotRoleARN = r.Form.Get("RoleArn") gotSessionName = r.Form.Get("RoleSessionName") + // With keep-alive disabled, Go's HTTP client sends Connection: close. + gotConnection = r.Header.Get("Connection") w.Header().Set("Content-Type", "text/xml") _, _ = w.Write([]byte(` @@ -205,6 +207,10 @@ func TestBuildBedrockCredentialsAssumeRole(t *testing.T) { require.Equal(t, "arn:aws:iam::123456789012:role/target", gotRoleARN) require.Equal(t, bedrockSessionName, gotSessionName) + // The STS client disables keep-alive so each AssumeRole opens a fresh + // connection; Go signals this with a Connection: close request header. + require.Equal(t, "close", gotConnection, + "STS client should disable keep-alives so each AssumeRole opens a fresh connection") } // TestBuildBedrockCredentialsAssumeRoleExternalID verifies that a configured