mirror of
https://github.com/coder/coder.git
synced 2026-10-16 14:48:56 +08:00
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:
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
|
||||
Reference in New Issue
Block a user