Thanks to visit codestin.com
Credit goes to github.com

Skip to content

fix(aibridge/provider): disable keep-alive on the STS assume-role client #26971

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
evgeniy-scherbina merged 2 commits into from
Jul 2, 2026
Merged

fix(aibridge/provider): disable keep-alive on the STS assume-role client #26971

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bd4b826
fix(aibridge/provider): disable keep-alive on the STS assume-role client
evgeniy-scherbina Jul 2, 2026
fa282e4
test: assert the STS client disables keep-alive
evgeniy-scherbina Jul 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 20 additions & 1 deletion aibridge/provider/bedrock.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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.
Comment on lines +82 to +93

@coder-agents-review coder-agents-review Bot Jul 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit [CRF-2] Comment is 12 lines; 4 carry the why-not-what and trap, the rest is narrative padding.

Gon proposed trimming to:

// Disable keep-alive: with connection reuse, AssumeRole pins to one STS
// endpoint and returns stale AccessDenied for minutes after a trust-policy
// change. A fresh connection recovers in seconds. AssumeRole runs once per
// credential-cache lifetime, so keep-alive saves nothing.

This cuts the "Observed:" investigation narrative, the AWS CLI verification detail, the "STS-internal reason is unconfirmed" speculation, and the scope sentence that repeats what the local stsClient variable already shows.

Orchestrator note: downgraded from Gon's P2. Three reviewers (Leorio, Knov, Mafu-san) independently praised this comment's structure, specifically its observation-inference separation and honest uncertainty disclosure. The trim is good advice; the original is not a defect. (Gon P2, downgraded by orchestrator)

🤖

@coder-agents-review coder-agents-review Bot Jul 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit [CRF-3] The commit subject names the mechanism ("disable keep-alive on the STS assume-role client") instead of the condition it fixes. A git blame reader needs the condition, not the knob. Consider: fix(aibridge/provider): prevent stale AccessDenied after STS trust-policy changes. The PR description is exemplary; the subject should carry the same intent. (Leorio)

🤖

stsClient := sts.NewFromConfig(base, func(o *sts.Options) {
o.HTTPClient = awshttp.NewBuildableClient().WithTransportOptions(func(t *http.Transport) {
t.DisableKeepAlives = true

@coder-agents-review coder-agents-review Bot Jul 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P3 [CRF-1] No test asserts that the STS client disables keep-alives, the entire point of this PR.

Delete the DisableKeepAlives line and every test still passes green. The fix could regress silently during a future refactor.

The existing TestBuildBedrockCredentialsAssumeRole mock handler already receives the HTTP request. When DisableKeepAlives is true, Go's HTTP client sends a Connection: close header. A single assertion in the mock handler proves the transport is wired correctly:

require.Equal(t, "close", r.Header.Get("Connection"),
    "STS client should disable keep-alives so each AssumeRole opens a fresh connection")

The thorough comment at lines 82-93 mitigates accidental removal, but comments don't prevent regressions. A regression here means multi-minute production recovery windows after trust-policy changes, a failure mode that only surfaces under specific conditions. (Bisky)

🤖

@evgeniy-scherbina evgeniy-scherbina Jul 2, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good suggestion, done: fa282e4

})
})
credsProvider = stscreds.NewAssumeRoleProvider(stsClient, cfg.RoleARN, func(o *stscreds.AssumeRoleOptions) {
o.RoleSessionName = bedrockSessionName
if cfg.ExternalID != "" {
o.ExternalID = aws.String(cfg.ExternalID)
Expand Down
8 changes: 7 additions & 1 deletion aibridge/provider/bedrock_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(`<AssumeRoleResponse xmlns="https://sts.amazonaws.com/doc/2011-06-15/">
Expand Down Expand Up @@ -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
Expand Down
Loading