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>
This commit is contained in:
DevCats
2026-05-28 11:08:00 -05:00
committed by GitHub
co-authored by copilot-swe-agent[bot]
parent 0d1340a430
commit 094fe971ad
2 changed files with 18 additions and 0 deletions
+6
View File
@@ -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.
@@ -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)