diff --git a/aibridge/intercept/bedrocksig/bedrocksig.go b/aibridge/intercept/bedrocksig/bedrocksig.go new file mode 100644 index 00000000000..f9f2e6b46d7 --- /dev/null +++ b/aibridge/intercept/bedrocksig/bedrocksig.go @@ -0,0 +1,109 @@ +// Package bedrocksig holds the shared AWS SigV4 signing and Bedrock Mantle +// routing helpers used by both the anthropic-go and openai-go interceptors. +// It depends only on stdlib and the AWS SDK so it can be imported from either +// SDK-specific interceptor package without pulling in the other SDK. +package bedrocksig + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "io" + "net/http" + "net/url" + "strings" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4" + "golang.org/x/xerrors" +) + +// SigningService is the AWS SigV4 service name for Bedrock Mantle. +const SigningService = "bedrock-mantle" + +// PRMUserAgent is Coder's AWS Partner Revenue Measurement (PRM) attribution +// marker for outbound Bedrock requests. It is appended to the User-Agent +// header so AWS can recognize the traffic as Coder-associated Bedrock usage. +const PRMUserAgent = "sdk-ua-app-id/APN_1.1%2Fpc_cdfmjwn8i6u8l9fwz8h82e4w3%24" + +// AppendPRMUserAgent appends the Coder PRM attribution marker to the request's +// User-Agent header. +func AppendPRMUserAgent(req *http.Request) { + if ua := req.Header.Get("User-Agent"); ua != "" { + req.Header.Set("User-Agent", ua+" "+PRMUserAgent) + } +} + +// SignMiddleware returns an stdlib HTTP middleware that SigV4-signs the request +// for the Bedrock Mantle service. It appends the PRM user-agent, reads and +// restores the body for hashing, then signs with the given credentials and +// region. Callers wrap it in their SDK-specific option.WithMiddleware adapter. +func SignMiddleware(creds aws.CredentialsProvider, region string) func(req *http.Request, next func(*http.Request) (*http.Response, error)) (*http.Response, error) { + signer := v4.NewSigner() + return func(req *http.Request, next func(*http.Request) (*http.Response, error)) (*http.Response, error) { + AppendPRMUserAgent(req) + + resolved, err := creds.Retrieve(req.Context()) + if err != nil { + return nil, xerrors.Errorf("mantle SigV4: resolve AWS credentials: %w", err) + } + + // SigV4 requires a payload hash, so read the body to hash it and then + // restore it for the downstream HTTP client to send. + var body []byte + if req.Body != nil { + body, err = io.ReadAll(req.Body) + if err != nil { + return nil, xerrors.Errorf("mantle SigV4: read request body: %w", err) + } + _ = req.Body.Close() + req.Body = io.NopCloser(bytes.NewReader(body)) + req.ContentLength = int64(len(body)) + } + + hash := sha256.Sum256(body) + if err := signer.SignHTTP(req.Context(), resolved, req, hex.EncodeToString(hash[:]), SigningService, region, time.Now()); err != nil { + return nil, xerrors.Errorf("mantle SigV4: sign request: %w", err) + } + 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 +// bare vendor segments. +func trimMantleSegments(path string) string { + path = strings.TrimSuffix(path, "/") + path = strings.TrimSuffix(path, "/anthropic/v1") + path = strings.TrimSuffix(path, "/openai/v1") + path = strings.TrimSuffix(path, "/anthropic") + path = strings.TrimSuffix(path, "/openai") + return path +} + +// BaseURLForModel resolves the upstream Mantle base URL for a given model. It +// parses rawBase, trims any existing vendor path segment, then re-appends the +// one the model requires: "anthropic." models use /anthropic, "openai." models +// use /openai/v1, and anything else (third-party) uses /v1 so the SDK's +// /chat/completions route hits the root /v1/chat/completions endpoint. +func BaseURLForModel(rawBase, model string) (string, error) { + u, err := url.Parse(rawBase) + if err != nil { + return "", xerrors.Errorf("mantle base URL: %w", err) + } + + u.Path = trimMantleSegments(u.Path) + + switch { + case strings.HasPrefix(model, "anthropic."): + u.Path += "/anthropic" + case strings.HasPrefix(model, "openai."): + u.Path += "/openai/v1" + default: + // Third-party models use the root /v1 path segment. + u.Path += "/v1" + } + return u.String(), nil +} diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index f8143b63d80..a3e98fda0d7 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -1,14 +1,10 @@ package messages import ( - "bytes" "context" - "crypto/sha256" - "encoding/hex" "encoding/json" "errors" "fmt" - "io" "math" "net/http" "strconv" @@ -21,7 +17,6 @@ import ( "github.com/anthropics/anthropic-sdk-go/shared" "github.com/anthropics/anthropic-sdk-go/shared/constant" "github.com/aws/aws-sdk-go-v2/aws" - v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4" "github.com/google/uuid" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/trace" @@ -32,6 +27,7 @@ import ( aibcontext "github.com/coder/coder/v2/aibridge/context" "github.com/coder/coder/v2/aibridge/intercept" "github.com/coder/coder/v2/aibridge/intercept/apidump" + "github.com/coder/coder/v2/aibridge/intercept/bedrocksig" "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/mcp" "github.com/coder/coder/v2/aibridge/recorder" @@ -69,24 +65,10 @@ var bedrockSupportedBetaFlags = map[string]bool{ "tool-examples-2025-10-29": true, } -// BedrockPRMUserAgent is Coder's AWS Partner Revenue Measurement (PRM) -// attribution marker for outbound Bedrock requests. -// -// It is appended to Bedrock User-Agent headers so AWS can recognize the -// traffic as Coder-associated Bedrock usage. -const BedrockPRMUserAgent = "sdk-ua-app-id/APN_1.1%2Fpc_cdfmjwn8i6u8l9fwz8h82e4w3%24" - -// bedrockMantleSigningService is the AWS SigV4 service name for mantle. -const bedrockMantleSigningService = "bedrock-mantle" - -func appendBedrockPRMUserAgent(req *http.Request) { - if ua := req.Header.Get("User-Agent"); ua != "" { - req.Header.Set("User-Agent", ua+" "+BedrockPRMUserAgent) - } -} - // BedrockRuntime carries everything a Bedrock-backed interception needs: the -// static Bedrock config plus the AWS credentials provider. +// static Bedrock config plus the AWS credentials provider. The messages +// interceptor reads the full config (Model and SmallFastModel for the +// InvokeModel remap, BaseURL, Region). type BedrockRuntime struct { Cfg aibconfig.AWSBedrock Creds aws.CredentialsProvider @@ -350,7 +332,7 @@ func (i *interceptionBase) withBedrockInvokeModelOptions(ctx context.Context) ([ var out []option.RequestOption out = append(out, option.WithMiddleware(func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) { - appendBedrockPRMUserAgent(req) + bedrocksig.AppendPRMUserAgent(req) return next(req) })) out = append(out, bedrock.WithConfig(awsCfg)) @@ -384,39 +366,12 @@ func (i *interceptionBase) withBedrockMantleOptions(ctx context.Context) ([]opti return nil, xerrors.Errorf("resolve AWS credentials: %w", err) } - signer := v4.NewSigner() var out []option.RequestOption out = append(out, option.WithBaseURL(cfg.BaseURL)) // Appended last so it runs innermost (right before the HTTP send) and signs // the request after all other headers are set. - out = append(out, option.WithMiddleware(func(req *http.Request, next option.MiddlewareNext) (*http.Response, error) { - appendBedrockPRMUserAgent(req) - - creds, err := i.bedrock.Creds.Retrieve(req.Context()) - if err != nil { - return nil, xerrors.Errorf("mantle SigV4: resolve AWS credentials: %w", err) - } - - // SigV4 requires a payload hash, so read the body to hash it and then - // restore it for the downstream HTTP client to send. - var body []byte - if req.Body != nil { - var err error - body, err = io.ReadAll(req.Body) - if err != nil { - return nil, xerrors.Errorf("mantle SigV4: read request body: %w", err) - } - _ = req.Body.Close() - req.Body = io.NopCloser(bytes.NewReader(body)) - req.ContentLength = int64(len(body)) - } - - hash := sha256.Sum256(body) - if err := signer.SignHTTP(req.Context(), creds, req, hex.EncodeToString(hash[:]), bedrockMantleSigningService, cfg.Region, time.Now()); err != nil { - return nil, xerrors.Errorf("mantle SigV4: sign request: %w", err) - } - return next(req) - })) + //nolint:bodyclose // bedrocksig.SignMiddleware reads and closes the request body in order to sign it. + out = append(out, option.WithMiddleware(bedrocksig.SignMiddleware(i.bedrock.Creds, cfg.Region))) return out, nil } diff --git a/aibridge/internal/integrationtest/bridge_internal_test.go b/aibridge/internal/integrationtest/bridge_internal_test.go index 60756131c9b..877d68ac4ba 100644 --- a/aibridge/internal/integrationtest/bridge_internal_test.go +++ b/aibridge/internal/integrationtest/bridge_internal_test.go @@ -36,7 +36,7 @@ import ( "github.com/coder/coder/v2/aibridge/config" "github.com/coder/coder/v2/aibridge/fixtures" "github.com/coder/coder/v2/aibridge/intercept" - "github.com/coder/coder/v2/aibridge/intercept/messages" + "github.com/coder/coder/v2/aibridge/intercept/bedrocksig" "github.com/coder/coder/v2/aibridge/internal/testutil" "github.com/coder/coder/v2/aibridge/mcp" "github.com/coder/coder/v2/aibridge/provider" @@ -378,7 +378,7 @@ func TestAWSBedrockIntegration(t *testing.T) { // Verify PRM attribution is appended to the User-Agent header. ua := received[0].Header.Get("User-Agent") - require.Contains(t, ua, messages.BedrockPRMUserAgent, + require.Contains(t, ua, bedrocksig.PRMUserAgent, "expected AWS PRM attribution in User-Agent header") interceptions := bridgeServer.Recorder.RecordedInterceptions() @@ -437,7 +437,7 @@ func TestAWSBedrockIntegration(t *testing.T) { "signature must be scoped to the bedrock-mantle service") require.Contains(t, received[0].Header.Get("User-Agent"), - messages.BedrockPRMUserAgent) + bedrocksig.PRMUserAgent) interceptions := bridgeServer.Recorder.RecordedInterceptions() require.Len(t, interceptions, 1)