[proxy] Bedrock cost-allocation metadata + per-provider metadata_disabled (#6791)

This commit is contained in:
Maycon Santos
2026-07-20 14:45:50 +02:00
committed by GitHub
parent 3fb26d458e
commit d64e9542eb
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")
}