From d6ab06e59fe633280c0d3c3ee23f8619fd6be1ea Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Thu, 16 Jul 2026 01:16:44 +0200 Subject: [PATCH] [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. --- .../modules/agentnetwork/catalog/catalog.go | 18 ++++ .../modules/agentnetwork/synthesizer.go | 10 ++- .../modules/agentnetwork/synthesizer_test.go | 88 +++++++++++++++++++ .../modules/agentnetwork/types/provider.go | 10 +++ .../agentnetwork/types/provider_test.go | 35 ++++++++ .../builtin/llm_identity_inject/factory.go | 5 ++ .../builtin/llm_identity_inject/middleware.go | 40 ++++++++- .../llm_identity_inject/middleware_test.go | 40 +++++++++ shared/management/http/api/openapi.yml | 9 ++ shared/management/http/api/types.gen.go | 6 ++ 10 files changed, 256 insertions(+), 5 deletions(-) diff --git a/management/internals/modules/agentnetwork/catalog/catalog.go b/management/internals/modules/agentnetwork/catalog/catalog.go index 962b30250..f82cffae6 100644 --- a/management/internals/modules/agentnetwork/catalog/catalog.go +++ b/management/internals/modules/agentnetwork/catalog/catalog.go @@ -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", diff --git a/management/internals/modules/agentnetwork/synthesizer.go b/management/internals/modules/agentnetwork/synthesizer.go index 74ac91845..95fe91773 100644 --- a/management/internals/modules/agentnetwork/synthesizer.go +++ b/management/internals/modules/agentnetwork/synthesizer.go @@ -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, } } diff --git a/management/internals/modules/agentnetwork/synthesizer_test.go b/management/internals/modules/agentnetwork/synthesizer_test.go index 9d55bddf1..206d1d12a 100644 --- a/management/internals/modules/agentnetwork/synthesizer_test.go +++ b/management/internals/modules/agentnetwork/synthesizer_test.go @@ -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 / diff --git a/management/internals/modules/agentnetwork/types/provider.go b/management/internals/modules/agentnetwork/types/provider.go index 2e3195481..b3287168e 100644 --- a/management/internals/modules/agentnetwork/types/provider.go +++ b/management/internals/modules/agentnetwork/types/provider.go @@ -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, } diff --git a/management/internals/modules/agentnetwork/types/provider_test.go b/management/internals/modules/agentnetwork/types/provider_test.go index 1195499e7..f9756bb8b 100644 --- a/management/internals/modules/agentnetwork/types/provider_test.go +++ b/management/internals/modules/agentnetwork/types/provider_test.go @@ -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") +} diff --git a/proxy/internal/middleware/builtin/llm_identity_inject/factory.go b/proxy/internal/middleware/builtin/llm_identity_inject/factory.go index 8594c392d..8a5314337 100644 --- a/proxy/internal/middleware/builtin/llm_identity_inject/factory.go +++ b/proxy/internal/middleware/builtin/llm_identity_inject/factory.go @@ -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 diff --git a/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go b/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go index ee3f1c20d..722588a15 100644 --- a/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go +++ b/proxy/internal/middleware/builtin/llm_identity_inject/middleware.go @@ -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 diff --git a/proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go b/proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go index aab1271d8..8ec0930b5 100644 --- a/proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go +++ b/proxy/internal/middleware/builtin/llm_identity_inject/middleware_test.go @@ -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) { diff --git a/shared/management/http/api/openapi.yml b/shared/management/http/api/openapi.yml index 529cd2225..47ca80a7c 100644 --- a/shared/management/http/api/openapi.yml +++ b/shared/management/http/api/openapi.yml @@ -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 diff --git a/shared/management/http/api/types.gen.go b/shared/management/http/api/types.gen.go index 4956f9a9b..a9e98cf84 100644 --- a/shared/management/http/api/types.gen.go +++ b/shared/management/http/api/types.gen.go @@ -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"`