From ab69fa2f0d7e13f5e4fa86fd95f20fed72088a0d Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 2 Jul 2026 14:58:17 -0400 Subject: [PATCH] fix(aibridge/provider): disable keep-alive on the STS assume-role client (#26971) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A Bedrock provider that assumes an IAM role kept failing with `AssumeRole` `AccessDenied` for several minutes after its target role's trust policy was changed, and only recovered on a gateway restart or a long wait. The request itself was correct: the AWS CLI, using the same identity and the same `ExternalId`/role/region, accepted the identical request immediately against the same endpoint. The difference is the connection. The Go SDK reuses a keep-alive connection for the STS client, so every `AssumeRole` rides one connection pinned to a single STS endpoint. After a trust-policy change, that connection kept returning `AccessDenied` for minutes while a fresh connection (the AWS CLI) accepted the identical request at once; it recovered only when the connection recycled or the process restarted. The exact STS-internal reason is unconfirmed (likely per-endpoint propagation of the change) — what is verified is that a fresh connection per call recovers promptly. Disable keep-alive on the STS client so each `AssumeRole` opens a fresh connection and a trust-policy update takes effect quickly. `AssumeRole` runs at most once per credential-cache lifetime, so keep-alive bought nothing here. The change is scoped to the STS client only; Bedrock model requests are signed by a separate client and keep their connection pooling. ## What the data proves | | CLI | Gateway | |--------------------|-------------------------------------------|------------------------------------------| | Identity / key | `bedrock-base-user-useless` / `AKIA…44NL` | same | | STS endpoint | `sts.us-east-2.amazonaws.com` | same | | Request params | `ExternalId=QL53…`, role, session, 900 | same | | Recovery after fix | 7 seconds (21:27:54) | ~4.5 minutes (21:32:17) | | Re-hitting AWS? | new call each time | yes — 77 fresh `AssumeRole`s, all denied | Same identity, params, and endpoint, concurrent — yet the gateway was denied for ~4.5 minutes while the CLI recovered in 7 seconds, and the gateway made a fresh `AssumeRole` on every request (so it was not caching a failure). The only difference was connection reuse. After disabling keep-alive, the same break/fix experiment brought gateway recovery down from ~4.5 minutes to ~7 seconds, in lockstep with the AWS CLI. --- aibridge/provider/bedrock.go | 21 ++++++++++++++++++++- aibridge/provider/bedrock_internal_test.go | 8 +++++++- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/aibridge/provider/bedrock.go b/aibridge/provider/bedrock.go index 013430f1d7..23f5b1db6a 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) diff --git a/aibridge/provider/bedrock_internal_test.go b/aibridge/provider/bedrock_internal_test.go index 1cb7764c89..e9827b7895 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