[proxy] forward Bedrock cost-allocation metadata + per-provider metadata_disabled

Stamp the caller's user and authorizing group onto AWS Bedrock requests via the
X-Amzn-Bedrock-Request-Metadata header (reusing the JSONMetadata identity-inject
shape with a Bedrock-charset sanitizer), so spend can be attributed in AWS Cost
Management. Add a per-provider metadata_disabled flag (default off) that
suppresses identity metadata injection while leaving catalog routing headers
intact.
This commit is contained in:
mlsmaycon
2026-07-16 01:16:44 +02:00
parent e1a24376ab
commit d6ab06e59f
10 changed files with 256 additions and 5 deletions

View File

@@ -197,6 +197,12 @@ type JSONMetadataInjection struct {
// enforces a 128-char limit per value; oversized values are
// truncated rather than failing the request. 0 disables the cap.
MaxValueLength int
// Sanitize, when true, replaces characters outside the destination's
// accepted set with '_' before emitting each value. AWS Bedrock's
// X-Amzn-Bedrock-Request-Metadata restricts values to a limited character
// class, so unsanitized group display names (e.g. containing spaces) would
// make Bedrock reject the request with 400.
Sanitize bool
}
// providers is the canonical list of supported Agent Network providers.
@@ -329,6 +335,18 @@ var providers = []Provider{
{ID: "amazon.nova-lite", Label: "Amazon Nova Lite (Bedrock)", InputPer1k: 0.00006, OutputPer1k: 0.00024, ContextWindow: 300000},
{ID: "amazon.nova-micro", Label: "Amazon Nova Micro (Bedrock)", InputPer1k: 0.000035, OutputPer1k: 0.00014, ContextWindow: 128000},
},
// Bedrock accepts a cost-allocation metadata header; stamp the caller's
// user + authorizing group so spend can be attributed in AWS Cost
// Management. Sanitized because Bedrock restricts the value character set.
IdentityInjection: &IdentityInjection{
JSONMetadata: &JSONMetadataInjection{
Header: "X-Amzn-Bedrock-Request-Metadata",
UserKey: "user",
GroupsKey: "group",
MaxValueLength: 256,
Sanitize: true,
},
},
},
{
ID: "vertex_ai_api",

View File

@@ -540,6 +540,7 @@ type identityInjectJSONMetadata struct {
UserKey string `json:"user_key,omitempty"`
GroupsKey string `json:"groups_key,omitempty"`
MaxValueLength int `json:"max_value_length,omitempty"`
Sanitize bool `json:"sanitize,omitempty"`
}
// buildIdentityInjectConfigJSON walks the enabled providers and emits
@@ -583,9 +584,11 @@ func buildIdentityInjectConfigJSON(providers []*types.Provider, groupIndex map[s
func buildIdentityInjectRule(p *types.Provider, entry catalog.Provider) (identityInjectProvider, bool) {
rule := identityInjectProvider{ProviderID: p.ID}
// Identity-stamping shape (one of HeaderPair / JSONMetadata). Skip the
// shape silently when the catalog entry doesn't declare one — extras
// can still apply, see below.
if entry.IdentityInjection != nil {
// shape silently when the catalog entry doesn't declare one, or when the
// operator disabled metadata for this provider — extras can still apply,
// see below. MetadataDisabled suppresses only the identity dimensions
// (user + authorizing group), not the catalog's routing ExtraHeaders.
if !p.MetadataDisabled && entry.IdentityInjection != nil {
switch {
case entry.IdentityInjection.HeaderPair != nil:
rule.HeaderPair = buildIdentityHeaderPair(p, entry.IdentityInjection.HeaderPair)
@@ -651,6 +654,7 @@ func buildIdentityJSONMetadata(p *types.Provider, jm *catalog.JSONMetadataInject
UserKey: userKey,
GroupsKey: groupsKey,
MaxValueLength: jm.MaxValueLength,
Sanitize: jm.Sanitize,
}
}

View File

@@ -698,6 +698,94 @@ func TestSynthesizeServices_IdentityInject_Portkey_NotCustomizable(t *testing.T)
"same fixed-schema guarantee for the groups dimension")
}
// TestSynthesizeServices_IdentityInject_Bedrock pins Bedrock's cost-allocation
// metadata: a JSONMetadata shape emitting X-Amzn-Bedrock-Request-Metadata with
// the reserved user/group keys, sanitized to Bedrock's accepted charset.
func TestSynthesizeServices_IdentityInject_Bedrock(t *testing.T) {
ctx := context.Background()
ctrl := gomock.NewController(t)
defer ctrl.Finish()
mockStore := store.NewMockStore(ctrl)
br := newSynthTestProvider()
br.ID = "prov-bedrock"
br.ProviderID = "bedrock_api"
br.UpstreamURL = "https://bedrock-runtime.us-east-1.amazonaws.com"
br.APIKey = "bedrock-bearer"
br.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC)
policy := newSynthTestPolicy(br.ID, "grp-eng", "")
policy.ID = "pol-bedrock"
expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(),
[]*types.Provider{br},
[]*types.Policy{policy},
[]*types.Guardrail{})
services, err := SynthesizeServices(ctx, mockStore, testAccountID)
require.NoError(t, err)
require.Len(t, services, 1)
var injectCfg identityInjectConfig
for _, m := range services[0].Targets[0].Options.Middlewares {
if m.ID == middlewareIDLLMIdentityInject {
require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg))
break
}
}
require.Len(t, injectCfg.Providers, 1)
entry := injectCfg.Providers[0]
require.NotNil(t, entry.JSONMetadata, "Bedrock uses the JSONMetadata shape for cost-allocation metadata")
assert.Nil(t, entry.HeaderPair, "shapes are mutually exclusive")
assert.Equal(t, "X-Amzn-Bedrock-Request-Metadata", entry.JSONMetadata.Header,
"the caller identity lands in Bedrock's cost-allocation metadata header")
assert.Equal(t, "user", entry.JSONMetadata.UserKey)
assert.Equal(t, "group", entry.JSONMetadata.GroupsKey)
assert.True(t, entry.JSONMetadata.Sanitize,
"Bedrock restricts the metadata value charset, so values must be sanitized")
}
// TestSynthesizeServices_MetadataDisabled_SuppressesInjection verifies the
// per-provider opt-out: a provider with MetadataDisabled set emits no
// identity-inject entry (Bedrock has no catalog ExtraHeaders, so the whole
// entry is dropped).
func TestSynthesizeServices_MetadataDisabled_SuppressesInjection(t *testing.T) {
ctx := context.Background()
ctrl := gomock.NewController(t)
defer ctrl.Finish()
mockStore := store.NewMockStore(ctrl)
br := newSynthTestProvider()
br.ID = "prov-bedrock"
br.ProviderID = "bedrock_api"
br.UpstreamURL = "https://bedrock-runtime.us-east-1.amazonaws.com"
br.APIKey = "bedrock-bearer"
br.MetadataDisabled = true
br.CreatedAt = time.Date(2026, 4, 2, 0, 0, 0, 0, time.UTC)
policy := newSynthTestPolicy(br.ID, "grp-eng", "")
policy.ID = "pol-bedrock"
expectSynthBaseInputs(mockStore, ctx, newSynthTestSettings(),
[]*types.Provider{br},
[]*types.Policy{policy},
[]*types.Guardrail{})
services, err := SynthesizeServices(ctx, mockStore, testAccountID)
require.NoError(t, err)
require.Len(t, services, 1)
var injectCfg identityInjectConfig
for _, m := range services[0].Targets[0].Options.Middlewares {
if m.ID == middlewareIDLLMIdentityInject {
require.NoError(t, json.Unmarshal(m.ConfigJSON, &injectCfg))
break
}
}
assert.Empty(t, injectCfg.Providers,
"metadata_disabled must drop the provider's identity-inject entry")
}
// TestSynthesizeServices_IdentityInject_Vercel pins Vercel AI
// Gateway's wiring: HeaderPair shape with fixed wire names dictated
// by Vercel's Custom Reporting API (ai-reporting-user /

View File

@@ -51,6 +51,12 @@ type Provider struct {
// private or self-signed certificate. The synthesiser propagates it into
// the router route so the proxy dials that provider's upstream insecurely.
SkipTLSVerification bool `gorm:"column:skip_tls_verification"`
// MetadataDisabled suppresses identity metadata injection for this provider.
// Metadata (the caller's user + authorizing group) is injected by default;
// when true the synthesiser omits the provider's identity-inject shape, so no
// user/group headers (e.g. Bedrock's X-Amzn-Bedrock-Request-Metadata) are
// stamped. Catalog ExtraHeaders (routing config) are unaffected.
MetadataDisabled bool `gorm:"column:metadata_disabled"`
// SessionPrivateKey + SessionPublicKey are the ed25519 keypair the
// synthesised reverse-proxy service uses to sign / verify session
// JWTs after a successful OIDC handshake. Generated once on
@@ -137,6 +143,9 @@ func (p *Provider) FromAPIRequest(req *api.AgentNetworkProviderRequest) {
if req.SkipTlsVerification != nil {
p.SkipTLSVerification = *req.SkipTlsVerification
}
if req.MetadataDisabled != nil {
p.MetadataDisabled = *req.MetadataDisabled
}
// Identity-header overrides for catalogs flagged Customizable.
// nil pointer = "field omitted on the wire" → leave the stored
// value untouched (per the openapi description). Empty string is
@@ -170,6 +179,7 @@ func (p *Provider) ToAPIResponse() *api.AgentNetworkProvider {
Models: models,
Enabled: p.Enabled,
SkipTlsVerification: p.SkipTLSVerification,
MetadataDisabled: p.MetadataDisabled,
CreatedAt: &created,
UpdatedAt: &updated,
}

View File

@@ -42,3 +42,38 @@ func TestProvider_SkipTLSVerification_RoundTrip(t *testing.T) {
assert.False(t, p.SkipTLSVerification, "explicit false must clear skip_tls_verification")
assert.False(t, p.ToAPIResponse().SkipTlsVerification, "response must reflect the cleared value")
}
// TestProvider_MetadataDisabled_RoundTrip covers the request→provider→response
// mapping of metadata_disabled, with the same update semantics: nil preserves,
// explicit false clears.
func TestProvider_MetadataDisabled_RoundTrip(t *testing.T) {
enable := true
disable := false
base := func() *api.AgentNetworkProviderRequest {
return &api.AgentNetworkProviderRequest{
ProviderId: "bedrock_api",
Name: "bedrock",
UpstreamUrl: "https://bedrock-runtime.us-east-1.amazonaws.com",
}
}
p := NewProvider("acc-1")
req := base()
req.MetadataDisabled = &enable
p.FromAPIRequest(req)
assert.True(t, p.MetadataDisabled, "create with metadata_disabled=true must set the field")
assert.True(t, p.ToAPIResponse().MetadataDisabled, "response must surface metadata_disabled")
// Omitting the field on update leaves the stored value untouched.
p.FromAPIRequest(base())
assert.True(t, p.MetadataDisabled, "omitting metadata_disabled on update must preserve it")
// Explicit false clears it (re-enables metadata).
req = base()
req.MetadataDisabled = &disable
p.FromAPIRequest(req)
assert.False(t, p.MetadataDisabled, "explicit false must clear metadata_disabled")
assert.False(t, p.ToAPIResponse().MetadataDisabled, "response must reflect the cleared value")
}

View File

@@ -64,6 +64,11 @@ type JSONMetadataRule struct {
UserKey string `json:"user_key,omitempty"`
GroupsKey string `json:"groups_key,omitempty"`
MaxValueLength int `json:"max_value_length,omitempty"`
// Sanitize replaces characters outside the destination provider's accepted
// set with '_' before emitting each value. AWS Bedrock's
// X-Amzn-Bedrock-Request-Metadata restricts values to [A-Za-z0-9 +-=._:/@];
// group display names with other characters would otherwise 400.
Sanitize bool `json:"sanitize,omitempty"`
}
// Config is the on-wire configuration accepted by the factory. An

View File

@@ -292,15 +292,21 @@ func applyJSONMetadata(rule *JSONMetadataRule, in *middleware.Input) *middleware
mutations := &middleware.Mutations{}
mutations.HeadersRemove = append(mutations.HeadersRemove, rule.Header)
emit := func(v string) string {
if rule.Sanitize {
v = sanitizeMetadataValue(v)
}
return truncate(v, rule.MaxValueLength)
}
payload := map[string]string{}
if rule.UserKey != "" {
if identity := identityFor(in); identity != "" {
payload[rule.UserKey] = truncate(identity, rule.MaxValueLength)
payload[rule.UserKey] = emit(identity)
}
}
if rule.GroupsKey != "" {
if csv := authorisingTagsCSV(in); csv != "" {
payload[rule.GroupsKey] = truncate(csv, rule.MaxValueLength)
payload[rule.GroupsKey] = emit(csv)
}
}
if len(payload) == 0 {
@@ -359,6 +365,36 @@ func truncate(s string, maxBytes int) string {
return s[:maxBytes]
}
// sanitizeMetadataValue replaces any character outside AWS Bedrock's accepted
// request-metadata class — letters, digits, space, and + - = . _ : / @ — with
// '_'. This keeps values (notably the groups CSV, whose commas are rejected, and
// group display names with arbitrary characters) from making Bedrock reject the
// request with 400. The result stays opaque to the gateway.
func sanitizeMetadataValue(s string) string {
var b strings.Builder
b.Grow(len(s))
for _, r := range s {
if metadataCharAllowed(r) {
b.WriteRune(r)
} else {
b.WriteByte('_')
}
}
return b.String()
}
func metadataCharAllowed(r rune) bool {
switch {
case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9':
return true
}
switch r {
case ' ', '+', '-', '=', '.', '_', ':', '/', '@':
return true
}
return false
}
// tagsIDsFromAuthorising reads llm_router's authorising-groups metadata
// (a CSV of group ids) and returns the parsed slice. Returns nil when
// the key is absent or empty so the caller can fall back to the full

View File

@@ -304,6 +304,46 @@ func TestInject_JSONMetadata_TruncatesValues(t *testing.T) {
"per-value byte length must be capped at MaxValueLength")
}
// TestInject_JSONMetadata_Sanitize pins the AWS-Bedrock sanitization path: when
// Sanitize is set, characters outside Bedrock's accepted metadata class
// (notably the groups CSV comma and arbitrary characters in group display
// names) are replaced with '_' so Bedrock doesn't reject the request. Allowed
// characters (letters, digits, spaces, and @ . _ : / + - =) pass through.
func TestInject_JSONMetadata_Sanitize(t *testing.T) {
rule := ProviderInjection{
ProviderID: portkeyProvider,
JSONMetadata: &JSONMetadataRule{
Header: "X-Amzn-Bedrock-Request-Metadata",
UserKey: "user",
GroupsKey: "group",
MaxValueLength: 256,
Sanitize: true,
},
}
mw := New(Config{Providers: []ProviderInjection{rule}})
in := newInput(portkeyProvider, "alice", []string{"g1", "g2"})
in.UserEmail = "alice@example.com"
// Group display names carry characters Bedrock rejects (comma, '#'); the CSV
// join adds another comma between the two groups.
in.UserGroupNames = []string{"Eng,Team", "Ops#1"}
out, err := mw.Invoke(context.Background(), in)
require.NoError(t, err)
require.NotNil(t, out.Mutations)
require.Len(t, out.Mutations.HeadersAdd, 1)
added := out.Mutations.HeadersAdd[0]
assert.Equal(t, "X-Amzn-Bedrock-Request-Metadata", added.Key,
"the Bedrock cost-allocation header carries the metadata JSON")
var payload map[string]string
require.NoError(t, json.Unmarshal([]byte(added.Value), &payload))
assert.Equal(t, "alice@example.com", payload["user"],
"'@' and '.' are in Bedrock's accepted set and must be preserved")
assert.NotContains(t, payload["group"], ",", "commas must be sanitized — Bedrock rejects them")
assert.NotContains(t, payload["group"], "#", "disallowed characters must be sanitized")
assert.Contains(t, payload["group"], "Eng", "allowed characters must be preserved")
}
// TestInject_JSONMetadata_EmptyIdentity_StripsButDoesNotAdd verifies the
// anti-spoof Remove still fires when there's nothing to stamp.
func TestInject_JSONMetadata_EmptyIdentity_StripsButDoesNotAdd(t *testing.T) {

View File

@@ -5164,6 +5164,10 @@ components:
type: boolean
description: Whether upstream TLS certificate verification is skipped when the proxy dials this provider's URL. Intended for self-hosted / internal gateways behind a private or self-signed certificate.
example: false
metadata_disabled:
type: boolean
description: Whether identity metadata injection is disabled for this provider. When enabled (the default), the proxy stamps the caller's user and authorizing group onto upstream requests as provider-specific metadata (e.g. AWS Bedrock's X-Amzn-Bedrock-Request-Metadata header). Set true to suppress it.
example: false
created_at:
type: string
format: date-time
@@ -5184,6 +5188,7 @@ components:
- models
- enabled
- skip_tls_verification
- metadata_disabled
- created_at
- updated_at
AgentNetworkProviderRequest:
@@ -5240,6 +5245,10 @@ components:
type: boolean
description: Skip upstream TLS certificate verification when the proxy dials this provider's URL. For self-hosted / internal gateways behind a private or self-signed certificate. Defaults to false. When omitted on update, the stored value is left unchanged.
example: false
metadata_disabled:
type: boolean
description: Disable identity metadata injection (the caller's user + authorizing group) for this provider. Defaults to false (metadata is injected). When omitted on update, the stored value is left unchanged.
example: false
required:
- provider_id
- name

View File

@@ -2227,6 +2227,9 @@ type AgentNetworkProvider struct {
// IdentityHeaderUserId Wire header name the proxy stamps with the caller's display identity (user email or peer name) when the catalog entry's HeaderPair is `customizable`. Empty disables stamping for this dimension. Ignored when the catalog entry has a fixed HeaderPair (e.g. LiteLLM, Portkey). Used today by Bifrost: typical values are `x-bf-lh-netbird_user_id` (always-on log metadata) or `x-bf-dim-netbird_user_id` (Prometheus / OTEL — requires the label to be pre-declared in the gateway's `client.prometheus_labels` config).
IdentityHeaderUserId *string `json:"identity_header_user_id,omitempty"`
// MetadataDisabled Whether identity metadata injection is disabled for this provider. When enabled (the default), the proxy stamps the caller's user and authorizing group onto upstream requests as provider-specific metadata (e.g. AWS Bedrock's X-Amzn-Bedrock-Request-Metadata header). Set true to suppress it.
MetadataDisabled bool `json:"metadata_disabled"`
// Models Models exposed through this endpoint, with the operator's per-1k input/output prices. Empty means all catalog models are allowed at catalog prices.
Models []AgentNetworkProviderModel `json:"models"`
@@ -2278,6 +2281,9 @@ type AgentNetworkProviderRequest struct {
// IdentityHeaderUserId Wire header name for the caller's display identity. See AgentNetworkProvider.identity_header_user_id. When omitted on a request, the stored value is left unchanged; pass an empty string explicitly to clear it (which disables stamping for this dimension).
IdentityHeaderUserId *string `json:"identity_header_user_id,omitempty"`
// MetadataDisabled Disable identity metadata injection (the caller's user + authorizing group) for this provider. Defaults to false (metadata is injected). When omitted on update, the stored value is left unchanged.
MetadataDisabled *bool `json:"metadata_disabled,omitempty"`
// Models Models exposed through this endpoint, with the operator's per-1k input/output prices. Empty means all catalog models are allowed at catalog prices.
Models *[]AgentNetworkProviderModel `json:"models,omitempty"`