From 094fe971ad9226582efbef39fdd4ea5636af3ceb Mon Sep 17 00:00:00 2001 From: DevCats Date: Thu, 28 May 2026 11:08:00 -0500 Subject: [PATCH] chore(aibridge): add AWS PRM user-agent attribution for Bedrock calls (#25221) Adds middleware in `withAWSBedrockOptions` that appends the AWS Partner Revenue Measurement (PRM) attribution string to the User-Agent header on every Bedrock API call made through AI Bridge. This is the AI Bridge counterpart to the Terraform provisioner change merged in #23138. Together, they ensure all AWS API calls made by Coder (both workspace infrastructure via Terraform and LLM inference via Bedrock) include PRM attribution. ## How it works - A middleware is added before `bedrock.WithConfig(awsCfg)` that reads the existing `User-Agent` header and appends `sdk-ua-app-id/APN_1.1%2Fpc_cdfmjwn8i6u8l9fwz8h82e4w3%24` - Only affects Bedrock calls; OpenAI and direct Anthropic API calls are unaffected - Uses `option.WithMiddleware` rather than `option.WithHeader` because the existing User-Agent (set by the Anthropic SDK) must be preserved and appended to, not replaced ## Tests - **Positive**: `TestAWSBedrockIntegration` verifies PRM attribution is present in the User-Agent on Bedrock requests - **Negative**: `TestAnthropicMessages` verifies PRM attribution is absent on non-Bedrock requests ## References - Companion Terraform provisioner PR: #23138 (merged) - Backport: #24052 (merged) - Preserve existing `AWS_SDK_UA_APP_ID`: #24606 (open) - Original `coder/aibridge` PR: https://github.com/coder/aibridge/pull/224 (superseded by this PR since aibridge was moved into coder/coder via #24190) - [AWS SDK Application ID docs](https://docs.aws.amazon.com/sdkref/latest/guide/feature-appid.html) - [AWS PRM Automated User Agent](https://prm.partner.aws.dev/automated-user-agent.html) (partner login required) > Generated with [Coder Agents](https://coder.com/agents) Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- aibridge/intercept/messages/base.go | 6 ++++++ .../internal/integrationtest/bridge_internal_test.go | 12 ++++++++++++ 2 files changed, 18 insertions(+) diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index e35e2a9726..1f1f49e744 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -329,6 +329,12 @@ func (*interceptionBase) withAWSBedrockOptions(ctx context.Context, cfg *aibconf } var out []option.RequestOption + out = append(out, option.WithMiddleware(func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) { + if ua := req.Header.Get("User-Agent"); ua != "" { + req.Header.Set("User-Agent", ua+" sdk-ua-app-id/APN_1.1%2Fpc_cdfmjwn8i6u8l9fwz8h82e4w3%24") + } + return next(req) + })) out = append(out, bedrock.WithConfig(awsCfg)) // If a custom base URL is set, override the default endpoint constructed by the bedrock middleware. diff --git a/aibridge/internal/integrationtest/bridge_internal_test.go b/aibridge/internal/integrationtest/bridge_internal_test.go index cfb40599d3..595d7159c6 100644 --- a/aibridge/internal/integrationtest/bridge_internal_test.go +++ b/aibridge/internal/integrationtest/bridge_internal_test.go @@ -131,6 +131,13 @@ func TestAnthropicMessages(t *testing.T) { require.Len(t, promptUsages, 1) assert.Equal(t, "read the foo file", promptUsages[0].Prompt) + // Verify PRM attribution is NOT present on non-Bedrock Anthropic requests. + received := upstream.receivedRequests() + require.Len(t, received, 1) + ua := received[0].Header.Get("User-Agent") + assert.NotContains(t, ua, "sdk-ua-app-id", + "PRM attribution should not be present on non-Bedrock requests") + bridgeServer.Recorder.VerifyAllInterceptionsEnded(t) }) } @@ -327,6 +334,11 @@ func TestAWSBedrockIntegration(t *testing.T) { require.False(t, gjson.GetBytes(received[0].Body, "model").Exists(), "model should be stripped from body") require.False(t, gjson.GetBytes(received[0].Body, "stream").Exists(), "stream should be stripped from body") + // Verify PRM attribution is appended to the User-Agent header. + ua := received[0].Header.Get("User-Agent") + require.Contains(t, ua, "sdk-ua-app-id/APN_1.1%2Fpc_cdfmjwn8i6u8l9fwz8h82e4w3%24", + "expected AWS PRM attribution in User-Agent header") + interceptions := bridgeServer.Recorder.RecordedInterceptions() require.Len(t, interceptions, 1) require.Equal(t, interceptions[0].Model, bedrockCfg.Model)