diff --git a/.github/workflows/agent-network-e2e.yml b/.github/workflows/agent-network-e2e.yml index 88b98293d..9501c5fba 100644 --- a/.github/workflows/agent-network-e2e.yml +++ b/.github/workflows/agent-network-e2e.yml @@ -12,6 +12,13 @@ on: AWS issues it. Leave empty for the Sonnet 4.6 default. required: false default: "" + test_pattern: + description: >- + Package pattern to run. Defaults to the whole suite; narrow it to one + package (e.g. ./e2e/agentnetwork/...) when a run only needs that + package's answer and not the sixteen minutes the container suite costs. + required: false + default: "./e2e/..." concurrency: group: ${{ github.workflow }}-${{ github.ref }} @@ -77,4 +84,8 @@ jobs: GOOGLE_VERTEX_PROJECT: ${{ secrets.E2E_GOOGLE_VERTEX_PROJECT }} GOOGLE_VERTEX_REGION: ${{ secrets.E2E_GOOGLE_VERTEX_REGION }} GOOGLE_VERTEX_MODEL: ${{ secrets.E2E_GOOGLE_VERTEX_MODEL }} - run: go test -tags e2e -timeout 40m -v ./e2e/... + # Read through an env var rather than interpolated into the run + # script: a dispatch input reaching a shell command directly is a + # script-injection seam, however trusted the dispatcher. + TEST_PATTERN: ${{ inputs.test_pattern || './e2e/...' }} + run: go test -tags e2e -timeout 40m -v "$TEST_PATTERN" diff --git a/e2e/agentnetwork/discovery_live_test.go b/e2e/agentnetwork/discovery_live_test.go index 321e751bf..22c9f31c2 100644 --- a/e2e/agentnetwork/discovery_live_test.go +++ b/e2e/agentnetwork/discovery_live_test.go @@ -169,17 +169,15 @@ func liveDiscoveryCases() []liveDiscoveryCase { // Bedrock lists inference profiles, not models: matchModelless routes // /inference-profiles to a Bedrock route and refuses /v1/models for one. // - // The request reaches AWS and AWS refuses it — bedrock-runtime answers - // , because ListInferenceProfiles is a CONTROL - // PLANE operation served by bedrock..amazonaws.com, not the runtime - // host. A provider record carries one upstream and it has to be the runtime - // host for InvokeModel to work, so no Bedrock record can serve a listing as - // the model stands today. + // The listing is served by the CONTROL PLANE (bedrock.), not the + // runtime host a provider record must point at for InvokeModel — the + // runtime host answers . The router now sends + // the listing, and only the listing, to the control plane, so this case + // asserts a real filtered listing rather than the 404 it used to get. // - // The mock upstream hides this entirely: it answers /inference-profiles on - // the same listener as everything else, so the routing test passes there - // while the real endpoint 404s. That is the whole reason this file exists, - // so the case is kept, asserting what actually happens. + // The mock upstream cannot show any of this: it answers + // /inference-profiles on the same listener as everything else, so a + // mock-based test passes whichever host the request went to. if k := os.Getenv("AWS_BEARER_TOKEN_BEDROCK"); k != "" { region := os.Getenv("AWS_REGION") if region == "" { @@ -192,9 +190,13 @@ func liveDiscoveryCases() []liveDiscoveryCase { cases = append(cases, liveDiscoveryCase{ name: "bedrock", catalogID: "bedrock_api", upstream: "https://bedrock-runtime." + region + ".amazonaws.com", apiKey: k, - path: "/inference-profiles", - models: []string{sharedllm.NormalizeAnthropicModel(strings.TrimPrefix(model, "global."))}, - outcome: outcomeUpstreamNoListing, + path: "/inference-profiles", + // Registered verbatim, as an operator would copy it from AWS: the + // region prefix is what makes the id invocable, and the listing + // returns ids in exactly this form. + models: []string{model}, + outcome: outcomeFiltered, + permitted: []string{model}, }) } @@ -323,13 +325,19 @@ func runLiveDiscoveryCase(t *testing.T, ctx context.Context, tc liveDiscoveryCas code, body := callUntil(t, func() (int, string, error) { return cl.Get(ctx, endpoint, proxyIP, tc.path, tc.headers) }, 200) - t.Logf("[discovery] %s GET %s -> %d; body: %s", tc.name, tc.path, code, truncate(body, 4000)) - require.Equal(t, 200, code, "%s discovery must be served; body: %s", tc.name, truncate(body, 2000)) + // Status only, not the body. A Bedrock listing embeds inference-profile + // ARNs carrying the 12-digit AWS account id, and these job logs are + // readable by anyone who can see the run. The ids line below is the finding + // anyway. The failure paths below are the same log: a listing that fails to + // arrive is an AWS refusal naming the resource it refused, and that name is + // an ARN carrying the same account id. + t.Logf("[discovery] %s GET %s -> %d", tc.name, tc.path, code) + require.Equal(t, 200, code, "%s discovery must be served; response was %s", tc.name, bodyShape(body)) ids, ok := listingIDs(body) require.Truef(t, ok, - "%s answered discovery with something other than a {\"data\":[{\"id\":…}]} listing, which the filter forwards untouched — the caller would get an unbounded picker; body: %s", - tc.name, truncate(body, 2000)) + "%s answered discovery with something other than a {\"data\":[{\"id\":…}]} listing, which the filter forwards untouched — the caller would get an unbounded picker; response was %s", + tc.name, bodyShape(body)) sort.Strings(ids) t.Logf("[discovery] %s: %d ids after filtering: %s", tc.name, len(ids), strings.Join(ids, ", ")) @@ -342,8 +350,11 @@ func runLiveDiscoveryCase(t *testing.T, ctx context.Context, tc liveDiscoveryCas } for _, id := range ids { _, direct := permitted[id] - _, normalised := permitted[sharedllm.NormalizeAnthropicModel(id)] - assert.Truef(t, direct || normalised, + _, dated := permitted[sharedllm.NormalizeAnthropicModel(id)] + // Bedrock ids carry a region prefix and version suffix the record may + // not repeat; the proxy's filter tries the same forms. + _, bedrock := permitted[sharedllm.NormalizeBedrockModel(id)] + assert.Truef(t, direct || dated || bedrock, "%s offered %q, which no policy on this route permits — every entry the picker shows must be a request the guardrail would allow", tc.name, id) } for _, hidden := range tc.wantHidden { @@ -362,24 +373,39 @@ func isProxyError(body string) bool { } // listingIDs pulls the model ids out of a listing response. ok is false when -// the body is not the {"data":[{"id":…}]} shape the filter recognises. +// the body is neither envelope the proxy's filter recognises — the two must +// stay in step, or this test reports "not a listing" for a response the proxy +// filtered perfectly well. func listingIDs(body string) ([]string, bool) { var doc struct { + // OpenAI's shape, which Anthropic adopted. Data []struct { ID string `json:"id"` } `json:"data"` + // Bedrock returns inference-profile summaries under a key of its own, + // with the id under a field of its own. + Summaries []struct { + ID string `json:"inferenceProfileId"` + } `json:"inferenceProfileSummaries"` } if err := json.Unmarshal([]byte(body), &doc); err != nil { return nil, false } - if doc.Data == nil { - return nil, false + switch { + case doc.Data != nil: + ids := make([]string, 0, len(doc.Data)) + for _, entry := range doc.Data { + ids = append(ids, entry.ID) + } + return ids, true + case doc.Summaries != nil: + ids := make([]string, 0, len(doc.Summaries)) + for _, entry := range doc.Summaries { + ids = append(ids, entry.ID) + } + return ids, true } - ids := make([]string, 0, len(doc.Data)) - for _, entry := range doc.Data { - ids = append(ids, entry.ID) - } - return ids, true + return nil, false } func caseNames(cases []liveDiscoveryCase) []string { @@ -390,6 +416,27 @@ func caseNames(cases []liveDiscoveryCase) []string { return names } +// bodyShape describes a response without quoting any of it: its size and the +// top-level keys it arrived under. That is what a discovery failure is +// diagnosed from — which envelope the vendor answered with — and it is all +// that may go in a message rendered into a public job log, because the values +// underneath can carry an ARN and its account id. +func bodyShape(body string) string { + var doc map[string]json.RawMessage + if err := json.Unmarshal([]byte(body), &doc); err != nil { + return strconv.Itoa(len(body)) + " bytes, not a JSON object" + } + keys := make([]string, 0, len(doc)) + for key := range doc { + keys = append(keys, key) + } + sort.Strings(keys) + if len(keys) == 0 { + return strconv.Itoa(len(body)) + " bytes, an empty JSON object" + } + return strconv.Itoa(len(body)) + " bytes, keyed by: " + strings.Join(keys, ", ") +} + // truncate bounds a logged response body. A live catalogue can run to tens of // kilobytes, and the useful part is the front. func truncate(s string, limit int) string { diff --git a/management/internals/modules/agentnetwork/modeldiscovery/discovery.go b/management/internals/modules/agentnetwork/modeldiscovery/discovery.go index 37401820c..253cc63b3 100644 --- a/management/internals/modules/agentnetwork/modeldiscovery/discovery.go +++ b/management/internals/modules/agentnetwork/modeldiscovery/discovery.go @@ -191,7 +191,7 @@ func (c *Client) discoveryURL(entry catalog.Provider, req Request) (string, erro // A provider record carries no region field: the region lives // inside the upstream host the operator already configured, so // read it back out rather than asking them for it twice. - region = regionFromUpstream(entry, req.UpstreamURL) + region = RegionFromUpstream(entry, req.UpstreamURL) } if region == "" { return "", fmt.Errorf("%w: %s discovery needs a region, and none could be read from the provider upstream", @@ -207,13 +207,13 @@ func (c *Client) discoveryURL(entry catalog.Provider, req Request) (string, erro return target.String(), nil } -// regionFromUpstream recovers the region an operator embedded in the provider +// RegionFromUpstream recovers the region an operator embedded in the provider // upstream, by matching it against the catalog's own host template. Bedrock's // template is "bedrock-runtime..amazonaws.com" and Vertex's is // "-aiplatform.googleapis.com", so the region is whatever sits between // the fixed halves. Returns empty when the upstream does not match the // template, which is the case for a custom or proxied endpoint. -func regionFromUpstream(entry catalog.Provider, upstreamURL string) string { +func RegionFromUpstream(entry catalog.Provider, upstreamURL string) string { prefix, suffix, found := strings.Cut(entry.DefaultHost, catalog.RegionPlaceholder) if !found { return "" diff --git a/management/internals/modules/agentnetwork/modeldiscovery/discovery_test.go b/management/internals/modules/agentnetwork/modeldiscovery/discovery_test.go index fba2c97d1..133bd5148 100644 --- a/management/internals/modules/agentnetwork/modeldiscovery/discovery_test.go +++ b/management/internals/modules/agentnetwork/modeldiscovery/discovery_test.go @@ -490,7 +490,43 @@ func TestRegionFromUpstream(t *testing.T) { {"bedrock regionless without scheme", bedrock, "bedrock-runtime.amazonaws.com", ""}, } { t.Run(tc.name, func(t *testing.T) { - assert.Equal(t, tc.want, regionFromUpstream(tc.entry, tc.upstream)) + assert.Equal(t, tc.want, RegionFromUpstream(tc.entry, tc.upstream)) }) } } + +// bedrockGeoListing carries profiles from geographies the original prefix list +// did not name. Every one reduces to a catalog key, so every one must arrive +// priced — an unstripped geography is what made a real account's listing come +// back almost entirely at zero. +const bedrockGeoListing = `{"inferenceProfileSummaries":[ + {"inferenceProfileId":"jp.anthropic.claude-sonnet-5-20260514-v1:0", + "inferenceProfileName":"JP Anthropic Claude Sonnet 5","status":"ACTIVE","type":"SYSTEM_DEFINED"}, + {"inferenceProfileId":"au.anthropic.claude-haiku-4-5-20251001-v1:0", + "inferenceProfileName":"AU Anthropic Claude Haiku 4.5","status":"ACTIVE","type":"SYSTEM_DEFINED"}, + {"inferenceProfileId":"us-gov.anthropic.claude-sonnet-5-20260514-v1:0", + "inferenceProfileName":"GovCloud Anthropic Claude Sonnet 5","status":"ACTIVE","type":"SYSTEM_DEFINED"} +]}` + +func TestBedrockProfilesFromAnyGeographyArrivePriced(t *testing.T) { + cl, _ := newStubClient(http.StatusOK, bedrockGeoListing) + + models, err := cl.Fetch(context.Background(), Request{ + CatalogID: "bedrock_api", + UpstreamURL: "https://bedrock-runtime.eu-central-1.amazonaws.com", + APIKey: "aws-token", + }) + require.NoError(t, err) + require.Len(t, models, 3) + + for _, m := range models { + assert.True(t, m.PricingKnown, "%s must resolve to a catalog rate", m.ID) + assert.Greater(t, m.InputPer1k, 0.0, "input rate for %s", m.ID) + assert.Greater(t, m.OutputPer1k, 0.0, "output rate for %s", m.ID) + assert.Greater(t, m.CacheReadPer1k, 0.0, "cache-read rate for %s", m.ID) + } + + // The wire id is preserved whatever the pricing key reduced to: it is the + // only form that works at invoke time. + assert.Equal(t, "jp.anthropic.claude-sonnet-5-20260514-v1:0", models[0].ID) +} diff --git a/management/internals/modules/agentnetwork/synthesizer.go b/management/internals/modules/agentnetwork/synthesizer.go index 76944698e..66a19acd9 100644 --- a/management/internals/modules/agentnetwork/synthesizer.go +++ b/management/internals/modules/agentnetwork/synthesizer.go @@ -10,6 +10,7 @@ import ( "strings" "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog" + "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/modeldiscovery" "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/sessionkey" @@ -380,6 +381,9 @@ type routerProviderRoute struct { // proxy dials this provider's upstream. For self-hosted / internal gateways // behind a private or self-signed certificate. SkipTLSVerify bool `json:"skip_tls_verify,omitempty"` + // DiscoveryHost, when set, is the host serving this provider's model + // listing, for a vendor that does not serve it from the inference host. + DiscoveryHost string `json:"discovery_host,omitempty"` } // indexProviderGroups walks the enabled policies and returns, per @@ -447,6 +451,9 @@ func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][] if err != nil { return nil, fmt.Errorf("router config for provider %s: %w", p.ID, err) } + // Lookup rather than assume: an unknown provider id yields the zero + // entry, which declares no discovery and so contributes nothing. + catalogEntry, _ := catalog.Lookup(p.ProviderID) headerName, headerValue, gcpSAKeyB64, err := providerAuthHeader(p) if err != nil { return nil, err @@ -466,6 +473,7 @@ func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][] Bedrock: catalog.IsBedrockPathStyle(p.ProviderID), GCPServiceAccountKeyB64: gcpSAKeyB64, SkipTLSVerify: p.SkipTLSVerification, + DiscoveryHost: discoveryHost(catalogEntry, p.UpstreamURL), }) } out, err := json.Marshal(cfg) @@ -475,6 +483,33 @@ func buildRouterConfigJSON(providers []*types.Provider, groupIndex map[string][] return out, nil } +// discoveryHost returns the host serving this provider's model listing when it +// differs from the inference host, and empty when the two are the same — which +// is true of every vendor but Bedrock, whose ListInferenceProfiles is a control +// plane operation on bedrock. while InvokeModel must go to +// bedrock-runtime.. One provider record therefore needs two hosts. +// +// The catalog declares the listing host; the region is recovered from the +// upstream the operator configured, since a provider record carries no region +// field. An upstream matching no catalog template yields empty rather than a +// guess: a proxied or self-hosted Bedrock endpoint may serve both from one +// place, and inventing a host would send the credential somewhere the operator +// never configured. +func discoveryHost(entry catalog.Provider, upstreamURL string) string { + if entry.Discovery == nil || entry.Discovery.Host == "" { + return "" + } + host := entry.Discovery.Host + if !strings.Contains(host, catalog.RegionPlaceholder) { + return host + } + region := modeldiscovery.RegionFromUpstream(entry, upstreamURL) + if region == "" { + return "" + } + return strings.ReplaceAll(host, catalog.RegionPlaceholder, region) +} + // providerVendor returns the parser surface ("openai", "anthropic", …) // the provider speaks, sourced from its catalog entry's ParserID. The // router uses it to keep a request the parser tagged with a vendor on a diff --git a/management/internals/modules/agentnetwork/synthesizer_pricing_test.go b/management/internals/modules/agentnetwork/synthesizer_pricing_test.go index 83961878a..e82f2ef05 100644 --- a/management/internals/modules/agentnetwork/synthesizer_pricing_test.go +++ b/management/internals/modules/agentnetwork/synthesizer_pricing_test.go @@ -103,3 +103,37 @@ func TestBuildCostMeterConfig_OrphanAndGatewayProviders(t *testing.T) { assert.NotContains(t, cfg.Pricing.Providers, "prov-litellm", "empty-models gateway needs no per-record entry") assert.NotEmpty(t, cfg.Pricing.Defaults["openai"], "defaults still ship so the gateway's catalog-model traffic is priced") } + +// TestBuildCostMeterConfig_BedrockGeographyOutsideTheOriginalFour is the +// accounting half of the geography bug. The docs tell operators to register a +// Bedrock id exactly as AWS issues it, region prefix included, and the cost +// meter keys its table by the normalized form. While the geography was matched +// against a list of four, a profile issued anywhere else kept its prefix, +// missed the catalog entry it was meant to inherit from, and billed with a +// zero entry underneath the operator's own rates — so every cache bucket +// metered free and a model priced only by catalog defaults metered at nothing +// at all. +func TestBuildCostMeterConfig_BedrockGeographyOutsideTheOriginalFour(t *testing.T) { + for _, geo := range []string{"jp", "au", "ca", "sa", "us-gov"} { + t.Run(geo, func(t *testing.T) { + bedrock := &types.Provider{ + ID: "prov-bedrock", + ProviderID: "bedrock_api", + Enabled: true, + Models: []types.ProviderModel{ + {ID: geo + ".anthropic.claude-sonnet-5-20260514-v1:0", InputPer1k: 0.003, OutputPer1k: 0.015}, + }, + } + raw, err := buildCostMeterConfigJSON([]*types.Provider{bedrock}, map[string][]string{"prov-bedrock": {"grp"}}) + require.NoError(t, err) + cfg := decodeCostMeterConfig(t, raw) + + e, ok := cfg.Pricing.Providers["prov-bedrock"]["anthropic.claude-sonnet-5"] + require.True(t, ok, "a %s profile must key by the same normalized id the parser emits", geo) + assert.InDelta(t, 0.0003, e.CacheReadPer1k, 1e-9, + "cache read must be inherited from the bedrock default entry, not left at zero") + assert.InDelta(t, 0.00375, e.CacheCreationPer1k, 1e-9, + "cache creation must be inherited from the bedrock default entry, not left at zero") + }) + } +} diff --git a/management/internals/modules/agentnetwork/synthesizer_test.go b/management/internals/modules/agentnetwork/synthesizer_test.go index 817129571..352d36646 100644 --- a/management/internals/modules/agentnetwork/synthesizer_test.go +++ b/management/internals/modules/agentnetwork/synthesizer_test.go @@ -10,6 +10,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/catalog" "github.com/netbirdio/netbird/management/internals/modules/agentnetwork/types" rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" "github.com/netbirdio/netbird/management/server/store" @@ -1245,3 +1246,57 @@ func TestSynthesizeServices_EmptyAPIKey_FailsClosed(t *testing.T) { require.Error(t, err, "synthesis must refuse a provider with no api key") assert.Contains(t, err.Error(), "no api key", "error must surface the missing credential") } + +// TestDiscoveryHost pins which providers get a separate listing host. Getting +// this wrong in either direction is costly: a missing host leaves Bedrock +// discovery 404ing at AWS, and a host on the wrong provider would send that +// provider's listing — and its credential — somewhere the operator never +// configured. +func TestDiscoveryHost(t *testing.T) { + entry := func(id string) catalog.Provider { + p, ok := catalog.Lookup(id) + require.True(t, ok, "catalog entry %s must exist", id) + return p + } + + for _, tc := range []struct { + name string + entry catalog.Provider + upstream string + want string + }{ + { + // ListInferenceProfiles is a control-plane operation; the runtime + // host answers for it. + name: "bedrock splits the listing off the runtime host", + entry: entry("bedrock_api"), upstream: "https://bedrock-runtime.eu-central-1.amazonaws.com", + want: "bedrock.eu-central-1.amazonaws.com", + }, + { + name: "bedrock in another region", + entry: entry("bedrock_api"), upstream: "https://bedrock-runtime.us-west-2.amazonaws.com", + want: "bedrock.us-west-2.amazonaws.com", + }, + { + // A proxied Bedrock endpoint may well serve both from one place, + // and there is no region to read back out of it. + name: "proxied bedrock upstream yields no discovery host", + entry: entry("bedrock_api"), upstream: "https://bedrock.internal.example.com", + want: "", + }, + { + name: "openai serves its listing from the same host", + entry: entry("openai_api"), upstream: "https://api.openai.com", + want: "", + }, + { + name: "vertex serves its listing from the same host", + entry: entry("vertex_ai_api"), upstream: "https://us-east5-aiplatform.googleapis.com", + want: "", + }, + } { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.want, discoveryHost(tc.entry, tc.upstream)) + }) + } +} diff --git a/proxy/internal/middleware/builtin/llm_router/bedrock_discovery_test.go b/proxy/internal/middleware/builtin/llm_router/bedrock_discovery_test.go new file mode 100644 index 000000000..d21e33c21 --- /dev/null +++ b/proxy/internal/middleware/builtin/llm_router/bedrock_discovery_test.go @@ -0,0 +1,175 @@ +package llm_router + +import ( + "context" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/netbirdio/netbird/proxy/internal/middleware" +) + +// bedrockRoute is a Bedrock provider whose listing lives on the control plane +// while inference goes to the runtime host — the split this file is about. +func bedrockRoute(models []string, policies []ModelPolicyRule) ProviderRoute { + return ProviderRoute{ + ID: "prov-bedrock", + Bedrock: true, + Models: models, + ModelPolicies: policies, + UpstreamScheme: "https", + UpstreamHost: "bedrock-runtime.eu-central-1.amazonaws.com", + DiscoveryHost: "bedrock.eu-central-1.amazonaws.com", + AuthHeaderName: "Authorization", + AuthHeaderValue: "Bearer aws-token", + AllowedGroupIDs: []string{defaultTestGroup}, + } +} + +func getInput(path string) *middleware.Input { + return &middleware.Input{ + Slot: middleware.SlotOnRequest, + Method: http.MethodGet, + URL: "https://endpoint.netbird.local" + path, + UserGroups: []string{defaultTestGroup}, + } +} + +// TestBedrockListingGoesToTheControlPlane is the whole point of DiscoveryHost. +// ListInferenceProfiles is not an operation bedrock-runtime implements — it +// answers — so a listing forwarded to the +// inference upstream can only 404, however well it is routed. +func TestBedrockListingGoesToTheControlPlane(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{bedrockRoute(nil, nil)}}) + + out, err := mw.Invoke(context.Background(), getInput("/inference-profiles")) + require.NoError(t, err) + require.Equal(t, middleware.DecisionAllow, out.Decision) + require.NotNil(t, out.Mutations) + require.NotNil(t, out.Mutations.RewriteUpstream) + + assert.Equal(t, "bedrock.eu-central-1.amazonaws.com", out.Mutations.RewriteUpstream.Host) +} + +// TestBedrockInferenceStillGoesToTheRuntimeHost is the other half: the +// redirect must apply to the listing alone. Sending an InvokeModel call to the +// control plane would break every Bedrock request in the account. +func TestBedrockInferenceStillGoesToTheRuntimeHost(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{bedrockRoute(nil, nil)}}) + + in := newInputWithModelAndURL("anthropic.claude-haiku-4-5", + "https://endpoint.netbird.local/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/invoke") + out, err := mw.Invoke(context.Background(), in) + require.NoError(t, err) + require.Equal(t, middleware.DecisionAllow, out.Decision) + require.NotNil(t, out.Mutations.RewriteUpstream) + + assert.Equal(t, "bedrock-runtime.eu-central-1.amazonaws.com", out.Mutations.RewriteUpstream.Host) +} + +// TestIsListingPath guards the narrower reading of "model-less". Both the +// upstream redirect and the policy bound key on this, and the warming probe +// must be excluded from both: it carries no listing to filter, and pointing it +// at the control plane would warm a pool the inference requests never use. +func TestIsListingPath(t *testing.T) { + for path, want := range map[string]bool{ + "/v1/models": true, + "/inference-profiles": true, + "/bedrock/inference-profiles": true, + "/api/hello": false, + "/v1/models/gpt-4o": false, // the per-model lookup, routed elsewhere + "/v1/chat/completions": false, + } { + t.Run(path, func(t *testing.T) { + assert.Equal(t, want, isListingPath(path)) + }) + } +} + +// TestBedrockListingIsBoundByPolicy covers the case that was previously +// unreachable: filtering keyed on /v1/models alone, so a Bedrock listing was +// routed but never narrowed to what the caller may use. +func TestBedrockListingIsBoundByPolicy(t *testing.T) { + route := bedrockRoute( + []string{"eu.anthropic.claude-haiku-4-5-20251001-v1:0", "eu.anthropic.claude-sonnet-4-6"}, + []ModelPolicyRule{{ + GroupIDs: []string{defaultTestGroup}, + // A guardrail allowlist names the catalog key, which is the form an + // operator picks in the UI — not the region-prefixed wire id the + // record registers. + Models: []string{"anthropic.claude-haiku-4-5"}, + }}, + ) + mw := New(Config{Providers: []ProviderRoute{route}}) + + out, err := mw.Invoke(context.Background(), getInput("/inference-profiles")) + require.NoError(t, err) + require.NotNil(t, out.Mutations.RewriteUpstream) + + // Exact-string intersection would find nothing here and bound the listing + // to empty, handing the caller a picker with no models on a provider that + // works perfectly well. + assert.Equal(t, []string{"eu.anthropic.claude-haiku-4-5-20251001-v1:0"}, + out.Mutations.RewriteUpstream.DiscoveryModels) +} + +// TestBedrockListingWithoutADiscoveryHostFallsThrough keeps a proxied or +// self-hosted Bedrock endpoint working: the synthesiser emits no discovery +// host for one, and the listing must then go to the configured upstream rather +// than nowhere. +func TestBedrockListingWithoutADiscoveryHostFallsThrough(t *testing.T) { + route := bedrockRoute(nil, nil) + route.UpstreamHost = "bedrock.internal.example.com" + route.DiscoveryHost = "" + mw := New(Config{Providers: []ProviderRoute{route}}) + + out, err := mw.Invoke(context.Background(), getInput("/inference-profiles")) + require.NoError(t, err) + require.NotNil(t, out.Mutations.RewriteUpstream) + + assert.Equal(t, "bedrock.internal.example.com", out.Mutations.RewriteUpstream.Host) +} + +// TestBedrockProfileDetailHonoursTheModelTable covers GetInferenceProfile, +// which the listing filter cannot help with: it answers for one profile with a +// single object, not a set, so nothing narrows it on the way back. Authorising +// it by provider type alone would let any caller with a Bedrock route read the +// full configuration of every profile in the account. +// +// Both registration spellings are exercised, because a record may carry the +// raw profile id AWS issues or the catalog key it reduces to. +func TestBedrockProfileDetailHonoursTheModelTable(t *testing.T) { + const permitted = "eu.anthropic.claude-sonnet-5-20260514-v1:0" + + for _, registered := range []string{permitted, "anthropic.claude-sonnet-5"} { + t.Run(registered, func(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{bedrockRoute([]string{registered}, nil)}}) + + out, err := mw.Invoke(context.Background(), getInput("/inference-profiles/"+permitted)) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision, + "a profile the record registers must still resolve") + + denied, err := mw.Invoke(context.Background(), + getInput("/inference-profiles/eu.anthropic.claude-opus-5-20260514-v1:0")) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionDeny, denied.Decision, + "a profile outside the record's models must not be readable") + }) + } +} + +// TestBedrockProfileListingStaysModelLess pins the other half: the listing +// names no profile, so it must not be judged against the model table. It is +// bounded by DiscoveryModels in the response instead, and denying it here +// would take model discovery away from exactly the records that enumerate +// their models. +func TestBedrockProfileListingStaysModelLess(t *testing.T) { + mw := New(Config{Providers: []ProviderRoute{bedrockRoute([]string{"anthropic.claude-sonnet-5"}, nil)}}) + + out, err := mw.Invoke(context.Background(), getInput("/inference-profiles")) + require.NoError(t, err) + assert.Equal(t, middleware.DecisionAllow, out.Decision) +} diff --git a/proxy/internal/middleware/builtin/llm_router/factory.go b/proxy/internal/middleware/builtin/llm_router/factory.go index ae3d44a40..81b8727f1 100644 --- a/proxy/internal/middleware/builtin/llm_router/factory.go +++ b/proxy/internal/middleware/builtin/llm_router/factory.go @@ -50,6 +50,13 @@ type ProviderRoute struct { // under different allowlists must not offer either group the other's // models. Empty means no policy restricts models on this route. ModelPolicies []ModelPolicyRule `json:"model_policies,omitempty"` + // DiscoveryHost, when set, is the host that serves this provider's model + // listing, for a vendor that does not serve it from the same host as + // inference. Bedrock is why it exists: ListInferenceProfiles is a control + // plane operation on bedrock., while InvokeModel must go to + // bedrock-runtime., so one record genuinely needs two hosts. + // Empty means the listing is served from UpstreamHost like everything else. + DiscoveryHost string `json:"discovery_host,omitempty"` // Vertex marks a Google Vertex AI provider. Vertex requests carry the // model in the URL path, so the router selects this route by path // (isVertexPath) and bypasses the model/vendor table entirely. diff --git a/proxy/internal/middleware/builtin/llm_router/middleware.go b/proxy/internal/middleware/builtin/llm_router/middleware.go index 01981666c..b8d4b001b 100644 --- a/proxy/internal/middleware/builtin/llm_router/middleware.go +++ b/proxy/internal/middleware/builtin/llm_router/middleware.go @@ -242,10 +242,16 @@ func (m *Middleware) routeModelless(reqPath, surface, method string, userGroups if _, hadPrefix := splitBedrockNamespace(reqPath); hadPrefix { stripBedrockNamespace(out) } - // What the caller may actually use bounds what the picker may offer: - // every entry outside it is a request the chain will deny a moment - // later. - if reqPath == modelListingPath && out.Mutations != nil && out.Mutations.RewriteUpstream != nil { + if isListingPath(reqPath) && out.Mutations != nil && out.Mutations.RewriteUpstream != nil { + // A vendor that serves its listing from somewhere other than its + // inference upstream is redirected here, and only for the listing + // — every other request still goes to the configured upstream. + if route.DiscoveryHost != "" { + out.Mutations.RewriteUpstream.Host = route.DiscoveryHost + } + // What the caller may actually use bounds what the picker may + // offer: every entry outside it is a request the chain will deny a + // moment later. if models, bounded := discoverableModels(route, userGroups); bounded { out.Mutations.RewriteUpstream.DiscoveryModels = models } @@ -310,6 +316,20 @@ func discoverableModels(route ProviderRoute, userGroups []string) ([]string, boo for _, m := range route.Models { if _, ok := permitted[m]; ok { intersection[m] = struct{}{} + continue + } + // The two sides are not always written the same way. A Bedrock record + // may register the raw inference-profile id an operator copied from + // AWS while a guardrail allowlist names the catalog key, and comparing + // those verbatim finds nothing — which would bound a correctly + // configured provider's listing down to empty. routeClaimsModel + // already normalises the candidate for exactly this reason, and the + // listing bound has to agree with it or the picker disagrees with what + // the guardrail will actually allow. + if route.Bedrock { + if _, ok := permitted[llm.NormalizeBedrockModel(m)]; ok { + intersection[m] = struct{}{} + } } } return sortedModels(intersection), true @@ -472,6 +492,14 @@ const connectionWarmPath = "/api/hello" // alone. const modelListingPath = "/v1/models" +// isListingPath reports whether reqPath asks for a MODEL LISTING, as opposed +// to the other model-less endpoints. Only a listing gets an upstream redirect +// and a policy bound: the connection-warming probe carries no model list to +// filter, and rewriting its host would send the warm-up to the wrong pool. +func isListingPath(reqPath string) bool { + return reqPath == modelListingPath || isBedrockModelLessPath(reqPath) +} + // isModelLessPath reports whether reqPath is a known non-inference endpoint // that legitimately carries no model at all: the model listing and the // connection-warming probe. These must route to an upstream rather than @@ -513,7 +541,30 @@ func modelDetailID(reqPath string) (string, bool) { // gateway that does serve the lookup get a working answer. func isBedrockModelLessPath(reqPath string) bool { native, _ := splitBedrockNamespace(reqPath) - return native == "/inference-profiles" || strings.HasPrefix(native, "/inference-profiles/") + return native == "/inference-profiles" || strings.HasPrefix(native, bedrockProfileDetailPrefix) +} + +// bedrockProfileDetailPrefix precedes the identifier in a GetInferenceProfile +// lookup, once any gateway namespace is off the front. +const bedrockProfileDetailPrefix = "/inference-profiles/" + +// bedrockProfileID returns the inference profile a "/inference-profiles/{id}" +// lookup names. The listing beside it names none, which is what separates the +// two: a listing is a set the response filter can bound, while this answers +// for one profile with a single object no filter inspects. +// +// The id arrives as AWS issues it — region prefix and version suffix included +// — because that is the only form that works at invoke time. +func bedrockProfileID(reqPath string) (string, bool) { + native, _ := splitBedrockNamespace(reqPath) + if !strings.HasPrefix(native, bedrockProfileDetailPrefix) { + return "", false + } + id := strings.TrimPrefix(native, bedrockProfileDetailPrefix) + if id == "" { + return "", false + } + return id, true } // isVertexPath reports whether reqPath is a Google Vertex AI publisher @@ -653,7 +704,23 @@ func (m *Middleware) matchModelless(reqPath, method string, userGroups []string) var eligible func(ProviderRoute) bool switch { case isBedrockModelLessPath(reqPath): - eligible = func(r ProviderRoute) bool { return r.Bedrock } + if profile, isDetail := bedrockProfileID(reqPath); isDetail { + // A detail lookup names one profile, so it is authorised like any + // other per-model request rather than by provider type alone. The + // listing beside it is bounded by DiscoveryModels on the way back, + // but this answers with a single object no filter inspects — so + // without the check here, a caller reads the full configuration of + // every profile in the account, including the ones its policy + // never named. + // + // The id is normalised first: a record may register the raw + // profile id or the catalog key it reduces to, and routeClaimsModel + // expects the normalised form an inference request would carry. + wanted := llm.NormalizeBedrockModel(profile) + eligible = func(r ProviderRoute) bool { return r.Bedrock && routeClaimsModel(r, wanted) } + } else { + eligible = func(r ProviderRoute) bool { return r.Bedrock } + } case isModelLessPath(reqPath): // Vertex/Bedrock are path-routed and don't serve OpenAI-style // model-listing endpoints; including them here could rewrite a diff --git a/proxy/internal/proxy/discovery_filter.go b/proxy/internal/proxy/discovery_filter.go index c9d606970..d5502e1f1 100644 --- a/proxy/internal/proxy/discovery_filter.go +++ b/proxy/internal/proxy/discovery_filter.go @@ -97,6 +97,18 @@ func isPlainJSONListing(resp *http.Response) bool { return strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "application/json") } +// listingEnvelopes maps a listing's wrapper key to the field naming the model +// id inside it. Vendors did not converge on one shape: OpenAI's is what +// Anthropic adopted, while Bedrock returns inference-profile summaries under a +// key of its own. A body matching none of these is forwarded untouched. +var listingEnvelopes = []struct { + key string + idField string +}{ + {"data", "id"}, + {"inferenceProfileSummaries", "inferenceProfileId"}, +} + // filterListingBody returns the listing with unauthorised entries removed. // ok is false when the body is not a listing shape, in which case the // caller must forward the original bytes. @@ -105,38 +117,41 @@ func filterListingBody(body []byte, permitted map[string]struct{}) ([]byte, bool if err := json.Unmarshal(body, &doc); err != nil { return nil, false } - raw, present := doc["data"] - if !present { - return nil, false - } - var entries []map[string]json.RawMessage - if err := json.Unmarshal(raw, &entries); err != nil { - return nil, false - } - - kept := make([]map[string]json.RawMessage, 0, len(entries)) - for _, entry := range entries { - if entryPermitted(entry, permitted) { - kept = append(kept, entry) + for _, envelope := range listingEnvelopes { + raw, present := doc[envelope.key] + if !present { + continue + } + var entries []map[string]json.RawMessage + if err := json.Unmarshal(raw, &entries); err != nil { + return nil, false } - } - encoded, err := json.Marshal(kept) - if err != nil { - return nil, false + kept := make([]map[string]json.RawMessage, 0, len(entries)) + for _, entry := range entries { + if entryPermitted(entry, envelope.idField, permitted) { + kept = append(kept, entry) + } + } + + encoded, err := json.Marshal(kept) + if err != nil { + return nil, false + } + doc[envelope.key] = encoded + out, err := json.Marshal(doc) + if err != nil { + return nil, false + } + return out, true } - doc["data"] = encoded - out, err := json.Marshal(doc) - if err != nil { - return nil, false - } - return out, true + return nil, false } // entryPermitted reports whether a listing entry names a model the policy // authorises, trying every form the same model is written in. -func entryPermitted(entry map[string]json.RawMessage, permitted map[string]struct{}) bool { - raw, ok := entry["id"] +func entryPermitted(entry map[string]json.RawMessage, idField string, permitted map[string]struct{}) bool { + raw, ok := entry[idField] if !ok { return false } @@ -184,6 +199,13 @@ func modelIDForms(id string) []string { return nil } forms := []string{id, sharedllm.NormalizeAnthropicModel(id)} + // A Bedrock listing returns region-prefixed, version-suffixed profile ids + // ("eu.anthropic.claude-haiku-4-5-20251001-v1:0") while the record may + // register the catalog key. Stripping to the key is a no-op for ids that + // carry neither, so this costs nothing on the other surfaces. + if bedrock := sharedllm.NormalizeBedrockModel(id); bedrock != id { + forms = append(forms, bedrock) + } if slash := strings.Index(id, "/"); slash > 0 { if _, ok := gatewayNamespaces[id[:slash]]; ok { tail := id[slash+1:] diff --git a/proxy/internal/proxy/discovery_filter_test.go b/proxy/internal/proxy/discovery_filter_test.go index 103eac594..fd5666345 100644 --- a/proxy/internal/proxy/discovery_filter_test.go +++ b/proxy/internal/proxy/discovery_filter_test.go @@ -233,3 +233,40 @@ func TestModelDiscoveryFilter_OversizedBodyKeepsUpstreamHeaders(t *testing.T) { assert.Equal(t, strconv.Itoa(len(body)), resp.Header.Get("Content-Length"), "the Content-Length header must not be rewritten to the truncated prefix") } + +// TestFilterBedrockInferenceProfiles covers the second listing envelope. AWS +// returns inference-profile summaries under a key of its own with an id field +// of its own, so a filter that only knew OpenAI's shape forwarded a Bedrock +// listing whole — offering every profile in the account regardless of policy. +func TestFilterBedrockInferenceProfiles(t *testing.T) { + body := []byte(`{"inferenceProfileSummaries":[ + {"inferenceProfileId":"eu.anthropic.claude-haiku-4-5-20251001-v1:0","status":"ACTIVE"}, + {"inferenceProfileId":"eu.anthropic.claude-sonnet-4-6","status":"ACTIVE"}, + {"inferenceProfileId":"global.cohere.embed-v4:0","status":"ACTIVE"} + ]}`) + + // The permitted set holds what the record registers. Here that is the + // catalog key, while the vendor answers with region-prefixed wire ids — + // the two must still line up. + permitted := map[string]struct{}{"anthropic.claude-haiku-4-5": {}} + + out, ok := filterListingBody(body, permitted) + require.True(t, ok, "a Bedrock listing must be recognised as filterable") + + var doc struct { + Summaries []struct { + ID string `json:"inferenceProfileId"` + } `json:"inferenceProfileSummaries"` + } + require.NoError(t, json.Unmarshal(out, &doc)) + require.Len(t, doc.Summaries, 1) + assert.Equal(t, "eu.anthropic.claude-haiku-4-5-20251001-v1:0", doc.Summaries[0].ID) +} + +// TestFilterLeavesUnknownEnvelopesAlone keeps the best-effort contract: a body +// the filter cannot parse must reach the client exactly as the upstream sent +// it, rather than being rewritten into something shorter and wrong. +func TestFilterLeavesUnknownEnvelopesAlone(t *testing.T) { + _, ok := filterListingBody([]byte(`{"models":[{"name":"something"}]}`), map[string]struct{}{}) + assert.False(t, ok) +} diff --git a/shared/llm/model.go b/shared/llm/model.go index 4fb631520..881097bda 100644 --- a/shared/llm/model.go +++ b/shared/llm/model.go @@ -10,9 +10,88 @@ import ( "strings" ) -// bedrockRegionPrefixes are the cross-region inference-profile prefixes that -// front a Bedrock model id (e.g. "eu.anthropic.claude-..."). -var bedrockRegionPrefixes = []string{"us.", "eu.", "apac.", "global."} +// bedrockVendorNamespaces are the vendor segments a Bedrock model id is +// published under. They identify the geography in front of a cross-region +// inference profile without knowing the geography: in +// "eu.anthropic.claude-...", what makes "eu" a geography is that "anthropic" +// follows it. +// +// A vendor missing from here is not fatal — bedrockGeographies covers the +// same id from the other side — but it is one of the two ways an id can go +// unrecognised, and the list needs a new entry whenever AWS onboards a +// vendor. A live listing found "global.xai.grok-4.6" days after this was +// first written. +var bedrockVendorNamespaces = map[string]struct{}{ + "ai21": {}, + "amazon": {}, + "anthropic": {}, + "cohere": {}, + "deepseek": {}, + "luma": {}, + "meta": {}, + "mistral": {}, + "openai": {}, + "qwen": {}, + "stability": {}, + "twelvelabs": {}, + "writer": {}, + "xai": {}, +} + +// bedrockGeographies are the geography segments AWS issues cross-region +// inference profiles under. They recognise a profile whose vendor we have +// never seen, which is the case bedrockVendorNamespaces alone gets wrong: +// "global.xai.grok-4.6" is a geography and a model whether or not "xai" is +// a name we know. +// +// Neither list is sufficient alone. A geography list on its own is what this +// file started with, and it aged badly — it held us, eu, apac and global, so +// every profile issued under jp, au, ca, sa or us-gov carried its prefix into +// the pricing key, matched no catalog entry, and reported the model unpriced. +// A vendor list on its own misses a new vendor under a known geography. +// Together, an id has to be new on both axes at once to go unrecognised. +var bedrockGeographies = map[string]struct{}{ + "apac": {}, + "au": {}, + "ca": {}, + "eu": {}, + "global": {}, + "jp": {}, + "sa": {}, + "us": {}, + "us-gov": {}, +} + +// stripBedrockGeography removes the cross-region inference-profile geography +// from a Bedrock model id, leaving the "." form the catalog and +// the pricing table key on. +// +// A leading segment counts as a geography when it is one we know, or when a +// known vendor follows it. Either alone is enough: the id has to be new on +// both axes before its geography survives. +// +// The segment has to be followed by two more, so "amazon.nova-pro" stays a +// vendor and a model rather than becoming a geography and a model — cutting +// its first segment would strip the vendor away. Over-stripping is the +// dangerous direction, because the result also decides which route may claim +// a model. +func stripBedrockGeography(modelID string) string { + geo, rest, found := strings.Cut(modelID, ".") + if !found || geo == "" { + return modelID + } + vendor, _, found := strings.Cut(rest, ".") + if !found { + return modelID + } + if _, ok := bedrockGeographies[geo]; ok { + return rest + } + if _, ok := bedrockVendorNamespaces[vendor]; ok { + return rest + } + return modelID +} // bedrockVersionSuffix matches the trailing "-vN[:N]" or "-YYYYMMDD-vN[:N]" // version/throughput suffix of a Bedrock model id. @@ -37,12 +116,7 @@ func NormalizeBedrockModel(modelID string) string { m = m[i+1:] } } - for _, p := range bedrockRegionPrefixes { - if strings.HasPrefix(m, p) { - m = m[len(p):] - break - } - } + m = stripBedrockGeography(m) return bedrockVersionSuffix.ReplaceAllString(m, "") } diff --git a/shared/llm/model_test.go b/shared/llm/model_test.go index 5ce2ff497..077a650fb 100644 --- a/shared/llm/model_test.go +++ b/shared/llm/model_test.go @@ -60,3 +60,61 @@ func TestNormalizeAnthropicModel(t *testing.T) { require.Equal(t, want, NormalizeAnthropicModel(in), "normalize %q", in) } } + +// TestNormalizeBedrockModel_GeographiesBeyondTheOriginalFour covers the bug +// that made this vendor-anchored: the geography used to be matched against a +// list of four, so a profile issued anywhere else kept its prefix, missed the +// catalog key it was supposed to match, and reported the model unpriced. +func TestNormalizeBedrockModel_GeographiesBeyondTheOriginalFour(t *testing.T) { + for _, geo := range []string{"us", "eu", "apac", "global", "jp", "au", "ca", "sa", "us-gov", "il", "mx"} { + t.Run(geo, func(t *testing.T) { + got := NormalizeBedrockModel(geo + ".anthropic.claude-sonnet-5-20260514-v1:0") + require.Equal(t, "anthropic.claude-sonnet-5", got, + "a cross-region profile must reduce to the catalog key whatever geography issued it") + }) + } +} + +// TestNormalizeBedrockModel_KeepsAVendorItCannotMistakeForAGeography pins the +// direction that must never break: a plain "." id has no +// geography, and cutting its first segment would strip the vendor away and +// hand the id to whichever route claims the bare model name. +func TestNormalizeBedrockModel_KeepsAVendorItCannotMistakeForAGeography(t *testing.T) { + cases := map[string]string{ + "amazon.nova-pro-v1:0": "amazon.nova-pro", + "anthropic.claude-sonnet-5-v1:0": "anthropic.claude-sonnet-5", + "meta.llama3-3-70b-instruct-v1:0": "meta.llama3-3-70b-instruct", + "cohere.command-r-plus-v1:0": "cohere.command-r-plus", + // Unknown on both axes: neither the leading segment nor the one + // after it is a name we hold, so the id is left exactly as it came. + "xx.unknownvendor.some-model-v1:0": "xx.unknownvendor.some-model", + "Qwen/Qwen2.5-0.5B-Instruct": "Qwen/Qwen2.5-0.5B-Instruct", + } + for in, want := range cases { + t.Run(in, func(t *testing.T) { + require.Equal(t, want, NormalizeBedrockModel(in)) + }) + } +} + +// TestNormalizeBedrockModel_RecognisesAnIdNewOnOneAxis covers what a live +// eu-central-1 listing returned days after the vendor list was written: +// "global.xai.grok-4.6", a vendor the list did not hold. Anchoring only on the +// vendor left the geography in the key, so the id matched no catalog entry and +// the model metered at zero. Each id below is unfamiliar on one axis and +// recognised through the other. +func TestNormalizeBedrockModel_RecognisesAnIdNewOnOneAxis(t *testing.T) { + cases := map[string]string{ + // Known geography, vendor we had never seen (the live case). + "global.xai.grok-4.6": "xai.grok-4.6", + "eu.xai.grok-4.6": "xai.grok-4.6", + // Known vendor, geography outside the list. + "il.anthropic.claude-sonnet-5-20260514-v1:0": "anthropic.claude-sonnet-5", + "mx.amazon.nova-2-lite-v1:0": "amazon.nova-2-lite", + } + for in, want := range cases { + t.Run(in, func(t *testing.T) { + require.Equal(t, want, NormalizeBedrockModel(in)) + }) + } +}