mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-09 00:11:28 +02:00
[proxy] Bedrock cost-allocation metadata + per-provider metadata_disabled (#6791)
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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 /
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user