fix: apply user-supplied AWS Bedrock API keys instead of dropping them (#28724)

## Bug

Under `/agents/settings/api-keys`, a personal API key saved for a
Bedrock provider never took effect. AI Gateway forwarded the saved
value, but the Bedrock interceptors re-signed every request with the
deployment's AWS credentials (SigV4), overwriting the `Authorization`
header. The key was silently dropped, so users were unknowingly still
billed to the administrator's AWS account with no way to bring their own
credentials.

## Fix

Apply the user's Bedrock key instead of discarding it. AWS Bedrock API
keys are bearer tokens (`Authorization: Bearer <key>`), so when a
personal key is present AI Gateway now sends it as a bearer token and
skips SigV4 signing. This covers every Bedrock route:

- **InvokeModel** (`/v1/messages`): applies the same wire transform as
the signed path (`/model/{model}/invoke[-with-response-stream]`,
`anthropic_version`, `anthropic-beta` → body, `model`/`stream`
stripping) but sets a bearer `Authorization` header instead of signing.
- **Mantle** (`/v1/messages`, `/v1/responses`, `/v1/chat/completions`):
passthrough with a bearer `Authorization` header instead of signing.

The bearer paths resolve no deployment AWS credentials, so BYOK works
even where the deployment has none. Requests without a personal key keep
the existing SigV4 behavior.

This also reverts the earlier interim guard that hid the key field and
rejected Bedrock key upserts: the field is shown again and the API
accepts the keys, because they now work.

<details>
<summary>Implementation notes</summary>

**Gateway (`aibridge`)**
- `messages.newMessagesService` detects BYOK on a Bedrock interception
and routes it through bearer-token option builders instead of the SigV4
path.
- `withBedrockInvokeModelBYOKOptions` /
`bedrockInvokeModelBearerMiddleware` reimplement the Anthropic SDK's
`bedrock.WithConfig` wire transform without signing, because that helper
couples the transform with SigV4 and would resolve deployment
credentials a BYOK request must not depend on.
- `bedrocksig.BearerMiddleware` is shared by the Messages, Responses,
and chat-completions interceptors for the Mantle passthrough; each
installs it in place of `SignMiddleware` when the credential is BYOK.
- Streaming decoding is unchanged: the Bedrock event-stream decoder is
registered by the SDK package's `init` and selected by response
content-type, independent of `bedrock.WithConfig`.

**coderd / SDK / frontend / docs**
- Removed the interim guard and `supports_user_api_key` capability flag;
restored the API-key field for Bedrock providers.
- Docs describe the Bedrock personal key as an AWS Bedrock API key
forwarded as a bearer token for every model family Bedrock serves.

**Tests**
- Unit: bearer middleware transform + auth for InvokeModel; shared
bearer middleware; option-builder validation.
- Integration: end-to-end through the bridge for InvokeModel, Mantle
Messages, Mantle Responses, and Mantle chat completions, asserting
`Authorization: Bearer <key>`, no `X-Amz-*` headers, correct upstream
path/model, no `X-Api-Key` leak, and that InvokeModel streaming decodes
real AWS binary event-stream frames.
- coderd: Bedrock key upserts succeed.

</details>

Closes
[CODAGT-699](https://linear.app/codercom/issue/CODAGT-699/bedrock-provider-settings-for-byok-only-expose-one-api-key).

---

*This PR was generated by Coder Agents on behalf of @f0ssel.*
This commit is contained in:
Garrett Delfosse
2026-09-25 18:15:47 -04:00
committed by GitHub
parent 99c42af5fc
commit b55f721f24
13 changed files with 824 additions and 23 deletions
+10
View File
@@ -86,6 +86,16 @@ func SignMiddleware(creds aws.CredentialsProvider, region, service string) func(
}
}
// BearerMiddleware authenticates a request with a user-supplied AWS Bedrock API
// key instead of SigV4. Callers install it in place of SignMiddleware so no
// deployment AWS credentials are resolved for the request.
func BearerMiddleware(token string) func(req *http.Request, next func(*http.Request) (*http.Response, error)) (*http.Response, error) {
return func(req *http.Request, next func(*http.Request) (*http.Response, error)) (*http.Response, error) {
req.Header.Set("Authorization", "Bearer "+token)
return next(req)
}
}
// trimMantleSegments removes any trailing Mantle vendor path segments from the
// parsed URL path so BaseURLForModel can re-append the correct one. Applied in
// longest-first order so /anthropic/v1 and /openai/v1 collapse before their
+31
View File
@@ -130,6 +130,37 @@ func TestAppendPRMUserAgent(t *testing.T) {
})
}
// TestBearerMiddleware verifies a user Bedrock API key is sent as a bearer token
// with the body and path untouched, and that no SigV4 headers are produced.
func TestBearerMiddleware(t *testing.T) {
t.Parallel()
const token = "bedrock-api-key" //nolint:gosec // G101: test-only fake credential.
body := `{"model":"anthropic.claude-x","max_tokens":1}`
req, err := http.NewRequestWithContext(context.Background(), http.MethodPost,
"https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages", strings.NewReader(body))
require.NoError(t, err)
var captured *http.Request
var capturedBody string
//nolint:bodyclose // next returns a synthetic no-op response body.
_, err = awssig.BearerMiddleware(token)(req, func(r *http.Request) (*http.Response, error) {
captured = r
b, readErr := io.ReadAll(r.Body)
require.NoError(t, readErr)
capturedBody = string(b)
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody}, nil
})
require.NoError(t, err)
require.NotNil(t, captured)
require.Equal(t, "Bearer "+token, captured.Header.Get("Authorization"))
require.Empty(t, captured.Header.Get("X-Amz-Date"))
require.Equal(t, "/anthropic/v1/messages", captured.URL.Path)
require.JSONEq(t, body, capturedBody)
}
func TestBaseURLForModel(t *testing.T) {
t.Parallel()
+13 -7
View File
@@ -37,8 +37,8 @@ type interceptionBase struct {
cred intercept.Credential
// bedrockMantle is nil for non-Bedrock providers. When set, upstream
// calls are SigV4-signed against the Bedrock Mantle endpoint instead of
// using a key pool or BYOK secret.
// calls target the Bedrock Mantle endpoint and are SigV4-signed, or
// bearer-authenticated when the request carries a user Bedrock API key.
bedrockMantle *awssig.MantleConfig
// clientHeaders are the original HTTP headers from the client request.
@@ -92,9 +92,9 @@ func (i *interceptionBase) newCompletionsService(ctx context.Context) openai.Cha
opts = append(opts, option.WithMiddleware(mw))
}
// Bedrock mantle: install the SigV4 signing middleware last so it runs
// innermost (right before the HTTP send) and signs the request after all
// other headers are set.
// Bedrock mantle: install auth last so it runs innermost (right before the
// HTTP send) after all other headers are set. A user Bedrock API key is
// sent as a bearer token; otherwise the request is SigV4-signed.
if i.bedrockMantle != nil {
// Bedrock traffic carries Coder's PRM attribution marker. This runs
// before SigV4 signing so the marker is covered by the signature.
@@ -102,8 +102,14 @@ func (i *interceptionBase) newCompletionsService(ctx context.Context) openai.Cha
awssig.AppendPRMUserAgent(req)
return next(req)
}))
//nolint:bodyclose // SignMiddleware reads and closes only the request body.
opts = append(opts, option.WithMiddleware(awssig.SignMiddleware(i.bedrockMantle.Creds, i.bedrockMantle.Region, awssig.ServiceBedrockMantle)))
if byok, ok := intercept.AsBYOK(i.cred); ok {
i.logger.Debug(ctx, "using byok auth", slog.F("key_hint", byok.Hint()))
//nolint:bodyclose // The middleware returns the upstream response for the SDK to close.
opts = append(opts, option.WithMiddleware(awssig.BearerMiddleware(byok.Secret)))
} else {
//nolint:bodyclose // SignMiddleware reads and closes only the request body.
opts = append(opts, option.WithMiddleware(awssig.SignMiddleware(i.bedrockMantle.Creds, i.bedrockMantle.Region, awssig.ServiceBedrockMantle)))
}
}
return openai.NewChatCompletionService(opts...)
@@ -7,6 +7,7 @@ import (
"testing"
"time"
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/google/uuid"
"github.com/openai/openai-go/v3"
"github.com/stretchr/testify/assert"
@@ -16,6 +17,7 @@ import (
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/aibridge/config"
"github.com/coder/coder/v2/aibridge/intercept"
"github.com/coder/coder/v2/aibridge/intercept/awssig"
"github.com/coder/coder/v2/aibridge/internal/testutil"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/recorder"
@@ -410,3 +412,71 @@ func TestWriteUpstreamError(t *testing.T) {
})
}
}
// TestNewCompletionsServiceBedrockAuth verifies the Bedrock mantle service uses
// the user's key as a bearer token when present and SigV4 otherwise.
func TestNewCompletionsServiceBedrockAuth(t *testing.T) {
t.Parallel()
const userKey = "user-bedrock-api-key" //nolint:gosec // G101: test-only fake credential.
tests := []struct {
name string
cred intercept.Credential
// check asserts on the Authorization header the upstream received.
check func(t *testing.T, header http.Header)
}{
{
name: "byok uses bearer token",
cred: intercept.BYOK{Secret: userKey, Header: intercept.AuthHeaderAuthorization},
check: func(t *testing.T, header http.Header) {
require.Equal(t, "Bearer "+userKey, header.Get("Authorization"))
require.Empty(t, header.Get("X-Amz-Date"))
},
},
{
name: "centralized uses sigv4",
cred: intercept.AWSSigV4{AccessKey: "AKID"},
check: func(t *testing.T, header http.Header) {
require.Contains(t, header.Get("Authorization"), "AWS4-HMAC-SHA256")
require.Contains(t, header.Get("Authorization"), "/bedrock-mantle/aws4_request")
require.NotEmpty(t, header.Get("X-Amz-Date"))
},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
var received http.Header
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
received = r.Header.Clone()
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"id":"cmpl_1","object":"chat.completion","choices":[]}`))
}))
t.Cleanup(upstream.Close)
base := &interceptionBase{
req: &ChatCompletionNewParamsWrapper{
ChatCompletionNewParams: openai.ChatCompletionNewParams{Model: "openai.gpt-x"},
},
cred: tc.cred,
bedrockMantle: &awssig.MantleConfig{
BaseURL: upstream.URL,
Region: "us-east-1",
Creds: credentials.NewStaticCredentialsProvider("AKID", "secret", ""),
},
logger: slog.Make(),
}
svc := base.newCompletionsService(context.Background())
_, err := svc.New(context.Background(), base.req.ChatCompletionNewParams)
require.NoError(t, err)
require.NotNil(t, received)
tc.check(t, received)
require.Contains(t, received.Get("User-Agent"), awssig.PRMUserAgent)
})
}
}
+137 -7
View File
@@ -1,12 +1,15 @@
package messages
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"math"
"net/http"
"net/url"
"strconv"
"strings"
"time"
@@ -18,6 +21,8 @@ import (
"github.com/anthropics/anthropic-sdk-go/shared/constant"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/google/uuid"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/trace"
"golang.org/x/xerrors"
@@ -325,9 +330,12 @@ func isSmallFastModel(model string) bool {
// newMessagesService builds the SDK service used for upstream calls.
func (i *interceptionBase) newMessagesService(ctx context.Context, opts ...option.RequestOption) (anthropic.MessageService, error) {
byok, isBYOK := intercept.AsBYOK(i.cred)
bedrockBYOK := isBYOK && i.bedrock != nil
// Only BYOK sets its credential here. Centralized keys are injected
// per-attempt in the failover loop.
if byok, ok := intercept.AsBYOK(i.cred); ok {
if isBYOK && i.bedrock == nil {
i.logger.Debug(ctx, "using byok auth",
slog.F("auth_header", byok.Header), slog.F("key_hint", byok.Hint()),
)
@@ -362,9 +370,17 @@ func (i *interceptionBase) newMessagesService(ctx context.Context, opts ...optio
const bedrockCredentialResolutionTimeout = 30 * time.Second
if i.isBedrockInvokeModel() {
ctx, cancel := context.WithTimeout(ctx, bedrockCredentialResolutionTimeout)
defer cancel()
bedrockOpts, err := i.withBedrockInvokeModelOptions(ctx)
var (
bedrockOpts []option.RequestOption
err error
)
if bedrockBYOK {
bedrockOpts, err = i.withBedrockInvokeModelBYOKOptions(byok.Secret)
} else {
tctx, cancel := context.WithTimeout(ctx, bedrockCredentialResolutionTimeout)
defer cancel()
bedrockOpts, err = i.withBedrockInvokeModelOptions(tctx)
}
if err != nil {
return anthropic.MessageService{}, err
}
@@ -373,9 +389,17 @@ func (i *interceptionBase) newMessagesService(ctx context.Context, opts ...optio
}
if i.isBedrockMantle() {
ctx, cancel := context.WithTimeout(ctx, bedrockCredentialResolutionTimeout)
defer cancel()
bedrockOpts, err := i.withBedrockMantleOptions(ctx)
var (
bedrockOpts []option.RequestOption
err error
)
if bedrockBYOK {
bedrockOpts, err = i.withBedrockMantleBYOKOptions(byok.Secret)
} else {
tctx, cancel := context.WithTimeout(ctx, bedrockCredentialResolutionTimeout)
defer cancel()
bedrockOpts, err = i.withBedrockMantleOptions(tctx)
}
if err != nil {
return anthropic.MessageService{}, err
}
@@ -503,6 +527,112 @@ func (i *interceptionBase) withBedrockMantleOptions(ctx context.Context) ([]opti
return append(out, signed...), nil
}
// withBedrockInvokeModelBYOKOptions authenticates the InvokeModel protocol with
// a user Bedrock API key (bearer token), resolving no deployment credentials.
func (i *interceptionBase) withBedrockInvokeModelBYOKOptions(token string) ([]option.RequestOption, error) {
if i.bedrock == nil {
return nil, xerrors.New("nil bedrock runtime")
}
cfg := i.bedrock.Cfg
if err := cfg.Validate(); err != nil {
return nil, xerrors.Errorf("bedrock invoke-model config: %w", err)
}
baseURL := fmt.Sprintf("https://bedrock-runtime.%s.amazonaws.com", cfg.Region)
if cfg.BaseURL != "" {
baseURL = cfg.BaseURL
}
var out []option.RequestOption
out = append(out, option.WithBaseURL(baseURL))
//nolint:bodyclose // The middleware returns the upstream response for the SDK to close.
out = append(out, option.WithMiddleware(bedrockInvokeModelBearerMiddleware(token)))
return out, nil
}
// withBedrockMantleBYOKOptions authenticates the mantle protocol with a user
// Bedrock API key (bearer token), resolving no deployment credentials.
func (i *interceptionBase) withBedrockMantleBYOKOptions(token string) ([]option.RequestOption, error) {
if i.bedrock == nil {
return nil, xerrors.New("nil bedrock runtime")
}
cfg := i.bedrock.Cfg
if err := cfg.Validate(); err != nil {
return nil, xerrors.Errorf("bedrock mantle config: %w", err)
}
// Bedrock traffic carries Coder's PRM attribution marker; the bearer token
// replaces SigV4 and resolves no deployment credentials.
var out []option.RequestOption
out = append(out, option.WithBaseURL(cfg.BaseURL))
out = append(out, option.WithMiddleware(func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) {
awssig.AppendPRMUserAgent(req)
return next(req)
}))
//nolint:bodyclose // The middleware returns the upstream response for the SDK to close.
out = append(out, option.WithMiddleware(awssig.BearerMiddleware(token)))
return out, nil
}
// bedrockInvokeModelBearerMiddleware authenticates an InvokeModel request with a
// user Bedrock API key (bearer token) instead of SigV4.
//
// It reimplements the SDK bedrock.WithConfig wire transform because that helper
// couples the transform with SigV4 signing and would resolve deployment AWS
// credentials, which a BYOK request must not depend on.
func bedrockInvokeModelBearerMiddleware(token string) func(*http.Request, option.MiddlewareNext) (*http.Response, error) {
return func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) {
awssig.AppendPRMUserAgent(req)
if req.Body != nil {
body, err := io.ReadAll(req.Body)
if err != nil {
return nil, xerrors.Errorf("bedrock byok: read request body: %w", err)
}
_ = req.Body.Close()
if !gjson.GetBytes(body, "anthropic_version").Exists() {
body, _ = sjson.SetBytes(body, "anthropic_version", bedrock.DefaultVersion)
}
// Bedrock expects beta flags in the body, not the header.
if betaHeader := req.Header.Values("anthropic-beta"); len(betaHeader) > 0 {
req.Header.Del("anthropic-beta")
body, err = sjson.SetBytes(body, "anthropic_beta", betaHeader)
if err != nil {
return nil, xerrors.Errorf("bedrock byok: set anthropic_beta: %w", err)
}
}
if req.Method == http.MethodPost && bedrock.DefaultEndpoints[req.URL.Path] {
model := gjson.GetBytes(body, "model").String()
stream := gjson.GetBytes(body, "stream").Bool()
body, _ = sjson.DeleteBytes(body, "model")
body, _ = sjson.DeleteBytes(body, "stream")
method := "invoke"
if stream {
method = "invoke-with-response-stream"
}
req.URL.Path = fmt.Sprintf("/model/%s/%s", model, method)
req.URL.RawPath = fmt.Sprintf("/model/%s/%s", url.QueryEscape(model), method)
}
reader := bytes.NewReader(body)
req.Body = io.NopCloser(reader)
req.GetBody = func() (io.ReadCloser, error) {
_, err := reader.Seek(0, 0)
return io.NopCloser(reader), err
}
req.ContentLength = int64(len(body))
}
req.Header.Set(intercept.AuthHeaderAuthorization, "Bearer "+token)
return next(req)
}
}
// augmentRequestForBedrockInvokeModel changes the model used for the request since AWS Bedrock doesn't support
// Anthropics' model names. It also converts adaptive thinking to enabled with a budget for models that
// don't support adaptive thinking natively, or enabled thinking to adaptive for models that only support
@@ -1,24 +1,30 @@
package messages
import (
"bytes"
"context"
"io"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/anthropics/anthropic-sdk-go"
"github.com/anthropics/anthropic-sdk-go/bedrock"
"github.com/anthropics/anthropic-sdk-go/option"
"github.com/anthropics/anthropic-sdk-go/shared/constant"
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
"golang.org/x/xerrors"
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/aibridge/config"
"github.com/coder/coder/v2/aibridge/intercept"
"github.com/coder/coder/v2/aibridge/intercept/awssig"
"github.com/coder/coder/v2/aibridge/internal/testutil"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/mcp"
@@ -1519,3 +1525,134 @@ func TestRecordTokenUsage(t *testing.T) {
})
}
}
// capturedRequest holds the request seen by a captureNext MiddlewareNext and a
// fully buffered copy of its body.
type capturedRequest struct {
req *http.Request
body []byte
}
// captureNext returns a MiddlewareNext that records the request it receives into
// c, then returns a minimal 200 response.
func captureNext(t *testing.T, c *capturedRequest) option.MiddlewareNext {
t.Helper()
return func(r *http.Request) (*http.Response, error) {
c.req = r
if r.Body != nil {
b, err := io.ReadAll(r.Body)
require.NoError(t, err)
c.body = b
}
return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(bytes.NewReader(nil))}, nil
}
}
// TestBedrockInvokeModelBearerMiddleware verifies that a user-supplied Bedrock
// API key authenticates via Authorization: Bearer, without SigV4 signing, while
// still applying the InvokeModel wire transform (path rewrite, anthropic_version
// injection, anthropic-beta header relocation, model/stream stripping).
func TestBedrockInvokeModelBearerMiddleware(t *testing.T) {
t.Parallel()
const token = "bedrock-api-key-abc123" //nolint:gosec // G101: test-only fake credential.
tests := []struct {
name string
stream bool
wantMethod string
}{
{name: "blocking", stream: false, wantMethod: "invoke"},
{name: "streaming", stream: true, wantMethod: "invoke-with-response-stream"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
body, err := sjson.SetBytes([]byte(`{"max_tokens":1}`), "model", "anthropic.claude-x")
require.NoError(t, err)
body, err = sjson.SetBytes(body, "stream", tt.stream)
require.NoError(t, err)
req, err := http.NewRequestWithContext(t.Context(), http.MethodPost,
"https://bedrock-runtime.us-east-1.amazonaws.com/v1/messages", bytes.NewReader(body))
require.NoError(t, err)
req.Header.Set("anthropic-beta", "interleaved-thinking-2025-05-14")
req.Header.Set("User-Agent", "test-agent")
var c capturedRequest
//nolint:bodyclose // captureNext returns a synthetic no-op response body.
_, err = bedrockInvokeModelBearerMiddleware(token)(req, captureNext(t, &c))
require.NoError(t, err)
captured, capturedBody := c.req, c.body
require.NotNil(t, captured)
// Bearer auth, not SigV4.
require.Equal(t, "Bearer "+token, captured.Header.Get("Authorization"))
require.NotContains(t, captured.Header.Get("Authorization"), "AWS4-HMAC-SHA256")
require.Empty(t, captured.Header.Get("X-Amz-Date"))
require.Empty(t, captured.Header.Get("X-Amz-Security-Token"))
// InvokeModel path rewrite.
require.Equal(t, "/model/anthropic.claude-x/"+tt.wantMethod, captured.URL.Path)
// Body transform: model & stream stripped, anthropic_version added,
// anthropic-beta relocated to the body.
require.False(t, gjson.GetBytes(capturedBody, "model").Exists())
require.False(t, gjson.GetBytes(capturedBody, "stream").Exists())
require.Equal(t, bedrock.DefaultVersion, gjson.GetBytes(capturedBody, "anthropic_version").String())
require.Equal(t, "interleaved-thinking-2025-05-14", gjson.GetBytes(capturedBody, "anthropic_beta.0").String())
require.Empty(t, captured.Header.Get("anthropic-beta"))
// PRM attribution appended.
require.Contains(t, captured.Header.Get("User-Agent"), awssig.PRMUserAgent)
})
}
}
// TestWithBedrockBYOKOptions verifies option assembly and validation for the
// BYOK (bearer token) Bedrock paths.
func TestWithBedrockBYOKOptions(t *testing.T) {
t.Parallel()
t.Run("InvokeModelOK", func(t *testing.T) {
t.Parallel()
base := &interceptionBase{bedrock: &BedrockRuntime{Cfg: config.AWSBedrock{
Region: "us-east-1",
Model: "anthropic.claude-x",
SmallFastModel: "anthropic.claude-haiku",
}}}
opts, err := base.withBedrockInvokeModelBYOKOptions("tok")
require.NoError(t, err)
require.NotEmpty(t, opts)
})
t.Run("MantleOK", func(t *testing.T) {
t.Parallel()
base := &interceptionBase{bedrock: &BedrockRuntime{Cfg: config.AWSBedrock{
Region: "us-east-1",
BaseURL: "https://bedrock-mantle.us-east-1.api.aws/anthropic",
Protocol: config.BedrockProtocolMantle,
}}}
opts, err := base.withBedrockMantleBYOKOptions("tok")
require.NoError(t, err)
require.NotEmpty(t, opts)
})
t.Run("NilRuntime", func(t *testing.T) {
t.Parallel()
base := &interceptionBase{}
_, err := base.withBedrockInvokeModelBYOKOptions("tok")
require.ErrorContains(t, err, "nil bedrock runtime")
_, err = base.withBedrockMantleBYOKOptions("tok")
require.ErrorContains(t, err, "nil bedrock runtime")
})
t.Run("InvalidConfig", func(t *testing.T) {
t.Parallel()
base := &interceptionBase{bedrock: &BedrockRuntime{Cfg: config.AWSBedrock{}}}
_, err := base.withBedrockInvokeModelBYOKOptions("tok")
require.Error(t, err)
})
}
+13 -7
View File
@@ -48,8 +48,8 @@ type responsesInterceptionBase struct {
cfg intercept.Config
cred intercept.Credential
// bedrockMantle is nil for non-Bedrock providers. When set, upstream
// calls are SigV4-signed against the Bedrock Mantle endpoint instead of
// using a key pool or BYOK secret.
// calls target the Bedrock Mantle endpoint and are SigV4-signed, or
// bearer-authenticated when the request carries a user Bedrock API key.
bedrockMantle *awssig.MantleConfig
// clientHeaders are the original HTTP headers from the client request.
@@ -118,9 +118,9 @@ func (i *responsesInterceptionBase) newResponsesService(ctx context.Context) res
opts = append(opts, option.WithMiddleware(mw))
}
// Bedrock mantle: install the SigV4 signing middleware last so it runs
// innermost (right before the HTTP send) and signs the request after all
// other headers are set.
// Bedrock mantle: install auth last so it runs innermost (right before the
// HTTP send) after all other headers are set. A user Bedrock API key is
// sent as a bearer token; otherwise the request is SigV4-signed.
if i.bedrockMantle != nil {
// Bedrock traffic carries Coder's PRM attribution marker. This runs
// before SigV4 signing so the marker is covered by the signature.
@@ -128,8 +128,14 @@ func (i *responsesInterceptionBase) newResponsesService(ctx context.Context) res
awssig.AppendPRMUserAgent(req)
return next(req)
}))
//nolint:bodyclose // SignMiddleware reads and closes only the request body.
opts = append(opts, option.WithMiddleware(awssig.SignMiddleware(i.bedrockMantle.Creds, i.bedrockMantle.Region, awssig.ServiceBedrockMantle)))
if byok, ok := intercept.AsBYOK(i.cred); ok {
i.logger.Debug(ctx, "using byok auth", slog.F("key_hint", byok.Hint()))
//nolint:bodyclose // The middleware returns the upstream response for the SDK to close.
opts = append(opts, option.WithMiddleware(awssig.BearerMiddleware(byok.Secret)))
} else {
//nolint:bodyclose // SignMiddleware reads and closes only the request body.
opts = append(opts, option.WithMiddleware(awssig.SignMiddleware(i.bedrockMantle.Creds, i.bedrockMantle.Region, awssig.ServiceBedrockMantle)))
}
}
return responses.NewResponseService(opts...)
@@ -7,6 +7,7 @@ import (
"testing"
"time"
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/google/uuid"
"github.com/openai/openai-go/v3"
oairesponses "github.com/openai/openai-go/v3/responses"
@@ -17,6 +18,7 @@ import (
"cdr.dev/slog/v3"
"github.com/coder/coder/v2/aibridge/config"
"github.com/coder/coder/v2/aibridge/intercept"
"github.com/coder/coder/v2/aibridge/intercept/awssig"
"github.com/coder/coder/v2/aibridge/internal/testutil"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/aibridge/recorder"
@@ -709,3 +711,72 @@ func TestWriteUpstreamError(t *testing.T) {
})
}
}
// TestNewResponsesServiceBedrockAuth verifies the Bedrock mantle service uses
// the user's key as a bearer token when present and SigV4 otherwise.
func TestNewResponsesServiceBedrockAuth(t *testing.T) {
t.Parallel()
const userKey = "user-bedrock-api-key" //nolint:gosec // G101: test-only fake credential.
tests := []struct {
name string
cred intercept.Credential
// check asserts on the Authorization header the upstream received.
check func(t *testing.T, header http.Header)
}{
{
name: "byok uses bearer token",
cred: intercept.BYOK{Secret: userKey, Header: intercept.AuthHeaderAuthorization},
check: func(t *testing.T, header http.Header) {
require.Equal(t, "Bearer "+userKey, header.Get("Authorization"))
require.Empty(t, header.Get("X-Amz-Date"))
},
},
{
name: "centralized uses sigv4",
cred: intercept.AWSSigV4{AccessKey: "AKID"},
check: func(t *testing.T, header http.Header) {
require.Contains(t, header.Get("Authorization"), "AWS4-HMAC-SHA256")
require.Contains(t, header.Get("Authorization"), "/bedrock-mantle/aws4_request")
require.NotEmpty(t, header.Get("X-Amz-Date"))
},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
var received http.Header
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
received = r.Header.Clone()
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"id":"resp_1","object":"response","output":[]}`))
}))
t.Cleanup(upstream.Close)
reqPayload, err := NewRequestPayload([]byte(`{"model":"openai.gpt-x","input":"hi"}`))
require.NoError(t, err)
base := &responsesInterceptionBase{
reqPayload: reqPayload,
cred: tc.cred,
bedrockMantle: &awssig.MantleConfig{
BaseURL: upstream.URL,
Region: "us-east-1",
Creds: credentials.NewStaticCredentialsProvider("AKID", "secret", ""),
},
logger: slog.Make(),
}
svc := base.newResponsesService(context.Background())
_, err = svc.New(context.Background(), oairesponses.ResponseNewParams{Model: "openai.gpt-x"})
require.NoError(t, err)
require.NotNil(t, received)
tc.check(t, received)
require.Contains(t, received.Get("User-Agent"), awssig.PRMUserAgent)
})
}
}
@@ -4,6 +4,7 @@ import (
"bytes"
"context"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"encoding/json"
"fmt"
@@ -20,6 +21,8 @@ import (
"github.com/anthropics/anthropic-sdk-go/packages/ssestream"
"github.com/anthropics/anthropic-sdk-go/shared/constant"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream/eventstreamapi"
v4signer "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
"github.com/google/uuid"
"github.com/openai/openai-go/v3"
@@ -522,6 +525,261 @@ func TestAWSBedrockIntegration(t *testing.T) {
}
})
// BYOK: a user-supplied Bedrock API key must authenticate the upstream
// request as a bearer token instead of being dropped in favor of the
// deployment's SigV4 signing.
t.Run("byok/invoke-model", func(t *testing.T) {
for _, streaming := range []bool{true, false} {
t.Run(fmt.Sprintf("streaming=%v", streaming), func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
const userKey = "user-bedrock-api-key"
fix := fixtures.Parse(t, fixtures.AntSingleBuiltinTool)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
bedrockCfg := &config.AWSBedrock{
Region: "us-west-2",
AccessKey: "test-access-key",
AccessKeySecret: "test-secret-key",
Model: "danthropic",
SmallFastModel: "danthropic-mini",
BaseURL: upstream.URL,
}
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL,
withCustomProvider(aibridgetest.NewAnthropicProvider(t, anthropicCfg(upstream.URL, apiKey), bedrockCfg)),
)
reqBody, err := sjson.SetBytes(fix.Request(), "stream", streaming)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, reqBody,
http.Header{"X-Api-Key": {userKey}})
require.NoError(t, err)
defer resp.Body.Close()
if streaming {
_, err = io.ReadAll(resp.Body)
require.NoError(t, err)
}
received := upstream.ReceivedRequests()
require.Len(t, received, 1)
// The user key authenticates as a bearer token; no SigV4.
require.Equal(t, "Bearer "+userKey, received[0].Header.Get("Authorization"))
require.NotContains(t, received[0].Header.Get("Authorization"), "AWS4-HMAC-SHA256")
require.Empty(t, received[0].Header.Get("X-Amz-Date"))
require.Empty(t, received[0].Header.Get("X-Api-Key"))
// InvokeModel wire transform still applied.
pathParts := strings.Split(received[0].Path, "/")
require.True(t, len(pathParts) >= 3 && pathParts[1] == "model", "unexpected path: %s", received[0].Path)
require.Equal(t, bedrockCfg.Model, pathParts[2])
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")
require.Contains(t, received[0].Header.Get("User-Agent"), awssig.PRMUserAgent)
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
})
// BYOK streaming must still decode Bedrock's binary event-stream frames.
// The bearer path drops bedrock.WithConfig, but the SDK's event-stream
// decoder is registered by the bedrock package's init and selected by the
// response content-type, so decoding is preserved.
t.Run("byok/invoke-model/eventstream", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
const userKey = "user-bedrock-api-key-eventstream"
var gotAuth atomic.Value
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotAuth.Store(r.Header.Get("Authorization"))
_ = writeBedrockEventStream(w,
`{"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"anthropic.claude","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":5,"output_tokens":1}}}`,
`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}`,
`{"type":"content_block_stop","index":0}`,
`{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}`,
`{"type":"message_stop"}`,
)
}))
t.Cleanup(upstream.Close)
bedrockCfg := &config.AWSBedrock{
Region: "us-west-2",
AccessKey: "test-access-key",
AccessKeySecret: "test-secret-key",
Model: "danthropic",
SmallFastModel: "danthropic-mini",
BaseURL: upstream.URL,
}
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL,
withCustomProvider(aibridgetest.NewAnthropicProvider(t, anthropicCfg(upstream.URL, apiKey), bedrockCfg)),
)
reqBody, err := sjson.SetBytes(fixtures.Parse(t, fixtures.AntSingleBuiltinTool).Request(), "stream", true)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, reqBody,
http.Header{"X-Api-Key": {userKey}})
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
// The re-emitted client SSE proves the binary frames were decoded.
sp := aibridge.NewSSEParser()
require.NoError(t, sp.Parse(resp.Body))
assert.Contains(t, sp.AllEvents(), "message_start")
assert.Contains(t, sp.AllEvents(), "message_stop")
require.Equal(t, "Bearer "+userKey, gotAuth.Load())
})
t.Run("byok/mantle", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
const userKey = "user-bedrock-api-key-mantle"
fix := fixtures.Parse(t, fixtures.AntSingleBuiltinTool)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
bedrockCfg := &config.AWSBedrock{
Region: "us-west-2",
AccessKey: "test-access-key",
AccessKeySecret: "test-secret-key",
BaseURL: upstream.URL + "/anthropic",
Protocol: config.BedrockProtocolMantle,
}
wantModel := gjson.GetBytes(fix.Request(), "model").String()
require.NotEmpty(t, wantModel)
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL,
withCustomProvider(aibridgetest.NewAnthropicProvider(t, anthropicCfg(upstream.URL, apiKey), bedrockCfg)),
)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, fix.Request(),
http.Header{"X-Api-Key": {userKey}})
require.NoError(t, err)
defer resp.Body.Close()
received := upstream.ReceivedRequests()
require.Len(t, received, 1)
// Native passthrough, bearer auth, no SigV4.
require.Equal(t, "/anthropic/v1/messages", received[0].Path)
require.Equal(t, "Bearer "+userKey, received[0].Header.Get("Authorization"))
require.NotContains(t, received[0].Header.Get("Authorization"), "AWS4-HMAC-SHA256")
require.Empty(t, received[0].Header.Get("X-Amz-Date"))
require.Empty(t, received[0].Header.Get("X-Api-Key"))
require.Equal(t, wantModel, gjson.GetBytes(received[0].Body, "model").String(), "model should be forwarded unchanged")
require.Contains(t, received[0].Header.Get("User-Agent"), awssig.PRMUserAgent)
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
// A user Bedrock API key must also authenticate the mantle OpenAI routes,
// whichever auth header the client used to supply it.
t.Run("byok/mantle/openai routes", func(t *testing.T) {
t.Parallel()
const userKey = "user-bedrock-api-key-openai"
cases := []struct {
name string
fixture []byte
model string
requestPath string
upstreamPath string
clientAuth http.Header
}{
{
name: "responses bearer",
fixture: fixtures.OaiResponsesBlockingSimple,
model: "openai.gpt-5.6-luna",
requestPath: "/bedrock/v1/responses",
upstreamPath: "/openai/v1/responses",
clientAuth: http.Header{"Authorization": {"Bearer " + userKey}},
},
{
name: "responses x-api-key",
fixture: fixtures.OaiResponsesBlockingSimple,
model: "openai.gpt-5.6-luna",
requestPath: "/bedrock/v1/responses",
upstreamPath: "/openai/v1/responses",
clientAuth: http.Header{"X-Api-Key": {userKey}},
},
{
name: "chat completions bearer",
fixture: fixtures.OaiChatSimple,
model: "mistral.ministral-3-3b-instruct",
requestPath: "/bedrock/v1/chat/completions",
upstreamPath: "/v1/chat/completions",
clientAuth: http.Header{"Authorization": {"Bearer " + userKey}},
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong)
t.Cleanup(cancel)
fix := fixtures.Parse(t, tc.fixture)
upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix))
bedrockCfg := config.AWSBedrock{
Region: "us-west-2",
AccessKey: "test-access-key",
AccessKeySecret: "test-secret-key",
BaseURL: upstream.URL,
Protocol: config.BedrockProtocolMantle,
}
bridgeServer := newBridgeTestServer(ctx, t, upstream.URL,
withCustomProvider(aibridgetest.NewBedrockProvider(t, config.Anthropic{
Name: config.ProviderBedrock,
BaseURL: upstream.URL,
}, bedrockCfg)),
)
reqBody, err := sjson.SetBytes(fix.Request(), "model", tc.model)
require.NoError(t, err)
resp, err := bridgeServer.makeRequest(t, http.MethodPost, tc.requestPath, reqBody, tc.clientAuth)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
_, err = io.Copy(io.Discard, resp.Body)
require.NoError(t, err)
received := upstream.ReceivedRequests()
require.Len(t, received, 1)
require.Equal(t, tc.upstreamPath, received[0].Path)
require.Equal(t, tc.model, gjson.GetBytes(received[0].Body, "model").String())
// Bearer auth with the user's key, no SigV4.
require.Equal(t, "Bearer "+userKey, received[0].Header.Get("Authorization"))
require.Empty(t, received[0].Header.Get("X-Amz-Date"))
require.Empty(t, received[0].Header.Get("X-Api-Key"))
require.Contains(t, received[0].Header.Get("User-Agent"), awssig.PRMUserAgent)
interceptions := bridgeServer.Recorder.RecordedInterceptions()
require.Len(t, interceptions, 1)
require.Equal(t, tc.model, interceptions[0].Model)
bridgeServer.Recorder.VerifyAllInterceptionsEnded(t)
})
}
})
// Tests that Bedrock-incompatible fields are stripped and adaptive thinking
// is handled correctly per model. Different Bedrock model names trigger
// different behavior for beta flag filtering and field stripping.
@@ -2782,3 +3040,35 @@ func TestTokenUsageRecordedWithoutMCPProxier(t *testing.T) {
})
}
}
// writeBedrockEventStream writes each JSON event as an AWS binary event-stream
// "chunk" frame, matching the wire format the Bedrock InvokeModel streaming
// endpoint returns.
func writeBedrockEventStream(w http.ResponseWriter, events ...string) error {
w.Header().Set("Content-Type", "application/vnd.amazon.eventstream")
w.WriteHeader(http.StatusOK)
encoder := eventstream.NewEncoder()
for _, event := range events {
payload, err := json.Marshal(map[string]string{
"bytes": base64.StdEncoding.EncodeToString([]byte(event)),
})
if err != nil {
return err
}
if err := encoder.Encode(w, eventstream.Message{
Headers: eventstream.Headers{
{Name: eventstreamapi.MessageTypeHeader, Value: eventstream.StringValue(eventstreamapi.EventMessageType)},
{Name: eventstreamapi.EventTypeHeader, Value: eventstream.StringValue("chunk")},
{Name: eventstreamapi.ContentTypeHeader, Value: eventstream.StringValue("application/json")},
},
Payload: payload,
}); err != nil {
return err
}
}
if flusher, ok := w.(http.Flusher); ok {
flusher.Flush()
}
return nil
}
+37
View File
@@ -3781,6 +3781,43 @@ func TestUserAIProviderKeys(t *testing.T) {
require.Equal(t, "AI provider is disabled.", sdkErr.Message)
})
t.Run("AcceptsBedrockProvider", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)
adminClient := newChatClient(t)
firstUser := coderdtest.CreateFirstUser(t, adminClient.Client)
memberClientRaw, _ := coderdtest.CreateAnotherUser(t, adminClient.Client, firstUser.OrganizationID)
memberClient := codersdk.NewExperimentalClient(memberClientRaw)
provider, err := adminClient.CreateAIProvider(ctx, codersdk.CreateAIProviderRequest{
Type: codersdk.AIProviderTypeBedrock,
Name: "test-bedrock-user-key-" + uuid.NewString(),
Enabled: true,
BaseURL: "https://bedrock-runtime.us-east-1.amazonaws.com",
Settings: codersdk.AIProviderSettings{
Bedrock: &codersdk.AIProviderBedrockSettings{
Region: "us-east-1",
Model: "anthropic.claude-3-5-sonnet",
SmallFastModel: "anthropic.claude-3-5-haiku",
},
},
})
require.NoError(t, err)
// Bedrock providers accept a user-supplied Bedrock API key, which the
// AI gateway forwards as a bearer token.
cfg, err := memberClient.UpsertUserAIProviderKey(ctx, "me", provider.ID, codersdk.CreateUserAIProviderKeyRequest{APIKey: "test-user-api-key"})
require.NoError(t, err)
require.True(t, cfg.HasUserAPIKey)
configs, err := memberClient.ListUserAIProviderKeyConfigs(ctx, "me")
require.NoError(t, err)
listed := findUserAIProviderKeyConfig(t, configs, provider.ID)
require.NotNil(t, listed)
require.True(t, listed.HasUserAPIKey)
})
t.Run("RejectsLargeAPIKey", func(t *testing.T) {
t.Parallel()
+5
View File
@@ -113,6 +113,11 @@ not by per-provider key policy flags. When BYOK is enabled, users can save a
personal API key for any enabled AI provider. When BYOK is disabled, saved user
keys are ignored and users cannot add or update personal keys.
For AWS Bedrock providers, the personal key is an AWS Bedrock API key. AI
Gateway forwards it as a bearer token, which authenticates that user's requests
instead of the deployment-managed AWS credentials (SigV4). This applies to
every model family Bedrock serves, including Anthropic and OpenAI models.
For each provider request, Coder selects credentials in this order:
1. If BYOK is enabled and the user has saved a personal key for the selected
+1 -1
View File
@@ -171,7 +171,7 @@ while allowing individual users to bring their own key.
Coder Agents requests routed through AI Gateway are in-process control plane requests, not external client requests that send their own AI Gateway bearer token.
Coder Agents use the same global BYOK setting.
When BYOK is enabled, users can save personal API keys for any enabled AI provider from the Agents settings page.
When BYOK is enabled, users can save personal API keys for any enabled AI provider from the Agents settings page. For AWS Bedrock providers, the personal key is an AWS Bedrock API key that AI Gateway forwards as a bearer token in place of the deployment-managed AWS credentials, for every model family Bedrock serves.
Refer to [Agents credential selection](../agents/models.md#credential-selection) for the Agents-specific behavior.
Visit individual [client pages](./clients/index.md) for configuration details.
+9 -1
View File
@@ -151,7 +151,15 @@ curl "${CODER_URL}/api/v2/ai-gateway/${PROVIDER_NAME}/v1/messages" \
#### AWS credentials
Do not attach API keys to a Bedrock provider.
Do not attach API keys to a Bedrock provider. The provider itself always
authenticates with AWS credentials.
When [BYOK](auth.md#bring-your-own-key-byok) is enabled, a user can save a
personal Amazon Bedrock API key under their
[Agents API keys](../agents/models.md#user-api-keys-byok). AI Gateway then
sends that key as a bearer token for the user's requests instead of signing
them with the deployment's AWS credentials. Requests from users without a
personal key continue to use the credentials below.
AI Gateway resolves AWS credentials one of three ways: