mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
co-authored by
copilot-swe-agent[bot]
parent
0d1340a430
commit
094fe971ad
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user