mirror of
https://github.com/netbirdio/netbird.git
synced 2026-08-27 10:01:28 +02:00
[proxy,management] Serve Bedrock model discovery from the control plane
A Bedrock provider could never answer a discovery request. The router
routed GET /inference-profiles to the record's upstream, which has to be
bedrock-runtime.<region> for InvokeModel to work, and that host does not
implement the operation — AWS answers <UnknownOperationException/>.
ListInferenceProfiles lives on the control plane at bedrock.<region>.
Give the route a discovery host, taken from the catalog's declaration
with the region read back out of the configured upstream, and send the
listing — and only the listing — there. Inference is untouched, and a
proxied or self-hosted Bedrock endpoint gets no discovery host at all
rather than a guessed one, since inventing a host would send the
operator's credential somewhere they never configured.
Two things had to follow for the listing to be usable once it arrives.
The response filter only understood OpenAI's {data:[{id:…}]}, so a
Bedrock listing was forwarded whole — offering every profile in the
account whatever the policy said. It now recognises the
inferenceProfileSummaries envelope, and matches a listing id against the
record's models after stripping the region prefix and version suffix, so
the two spellings of one model line up.
The policy bound had the same problem from the other side: it intersected
by exact string, so a record registering the raw profile id while a
guardrail names the catalog key intersected to nothing and would have
bounded a working provider's listing down to empty. routeClaimsModel
already normalises the candidate for this reason; the bound now agrees
with it.
The live discovery e2e flips from asserting the 404 to asserting a real
filtered listing. The mock upstream cannot cover 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.
This commit is contained in:
@@ -0,0 +1,133 @@
|
||||
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 <UnknownOperationException/> — 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)
|
||||
}
|
||||
@@ -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.<region>, while InvokeModel must go to
|
||||
// bedrock-runtime.<region>, 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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -162,6 +177,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.LastIndex(id, "/"); slash >= 0 {
|
||||
tail := id[slash+1:]
|
||||
forms = append(forms, tail, sharedllm.NormalizeAnthropicModel(tail))
|
||||
|
||||
@@ -215,3 +215,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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user